[ { "path": ".gitignore", "content": ".DS_Store\nxcuserdata/\n" }, { "path": "LICENSE", "content": " Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright 2022 Hugging Face SAS.\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "README.md", "content": "# SAM2 Studio\n\nThis is a Swift demo app for SAM 2 Core ML models.\n\n![UI Screenshot](screenshot.png)\n\nSAM 2 (Segment Anything in Images and Videos), is a collection of foundation models from FAIR that aim to solve promptable visual segmentation in images and videos. See the [SAM 2 paper](https://arxiv.org/abs/2408.00714) for more information.\n\n## Quick Start ⚡️\n\nDownload the compiled version [here!](https://huggingface.co/coreml-projects/sam-2-studio).\n\n## How to Use\n\nIf you prefer to compile it yourself or want to use a larger model, simply download the repo, compile with Xcode and run. \nThe app comes with the Small version of the model, but you can replace it with one of the supported models:\n\n- [SAM 2.1 Tiny](https://huggingface.co/apple/coreml-sam2.1-tiny)\n- [SAM 2.1 Small](https://huggingface.co/apple/coreml-sam2.1-small)\n- [SAM 2.1 Base](https://huggingface.co/apple/coreml-sam2.1-baseplus)\n- [SAM 2.1 Large](https://huggingface.co/apple/coreml-sam2.1-large)\n\nFor the older models, please check out the [Apple](https://huggingface.co/apple) organisation on HuggingFace.\n\nThis demo supports images, video support will be coming later.\n\n### Selecting Objects\n\n- You can select one or more _foreground_ points to choose objects in the image. Each additional point is interpreted as a _refinement_ of the previous mask.\n- Use a _background_ point to indicate an area to be removed from the current mask.\n- You can use a _box_ to select an approximate area that contains the object you're interested in.\n\n## Converting Models\n\nIf you want to use a fine-tuned model, you can convert it using [this fork of the SAM 2 repo](https://github.com/huggingface/segment-anything-2/tree/coreml-conversion). Please, let us know what you use it for!\n\n## Feedback and Contributions\n\nFeedback, issues and PRs are welcome! Please, feel free to [get in touch](https://github.com/huggingface/sam2-swiftui/issues/new).\n\n## Citation\n\nTo cite the SAM 2 paper, model, or software, please use the below:\n\n```\n@article{ravi2024sam2,\n title={SAM 2: Segment Anything in Images and Videos},\n author={Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and Hu, Ronghang and Ryali, Chaitanya and Ma, Tengyu and Khedr, Haitham and R{\\\"a}dle, Roman and Rolland, Chloe and Gustafson, Laura and Mintun, Eric and Pan, Junting and Alwala, Kalyan Vasudev and Carion, Nicolas and Wu, Chao-Yuan and Girshick, Ross and Doll{\\'a}r, Piotr and Feichtenhofer, Christoph},\n journal={arXiv preprint arXiv:2408.00714},\n url={https://arxiv.org/abs/2408.00714},\n year={2024}\n}\n```\n\n\n\n" }, { "path": "SAM2-Demo/Assets.xcassets/AccentColor.colorset/Contents.json", "content": "{\n \"colors\" : [\n {\n \"idiom\" : \"universal\"\n }\n ],\n \"info\" : {\n \"author\" : \"xcode\",\n \"version\" : 1\n }\n}\n" }, { "path": "SAM2-Demo/Assets.xcassets/AppIcon.appiconset/Contents.json", "content": "{\n \"images\" : [\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"1x\",\n \"size\" : \"16x16\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"2x\",\n \"size\" : \"16x16\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"1x\",\n \"size\" : \"32x32\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"2x\",\n \"size\" : \"32x32\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"1x\",\n \"size\" : \"128x128\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"2x\",\n \"size\" : \"128x128\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"1x\",\n \"size\" : \"256x256\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"2x\",\n \"size\" : \"256x256\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"1x\",\n \"size\" : \"512x512\"\n },\n {\n \"idiom\" : \"mac\",\n \"scale\" : \"2x\",\n \"size\" : \"512x512\"\n }\n ],\n \"info\" : {\n \"author\" : \"xcode\",\n \"version\" : 1\n }\n}\n" }, { "path": "SAM2-Demo/Assets.xcassets/Contents.json", "content": "{\n \"info\" : {\n \"author\" : \"xcode\",\n \"version\" : 1\n }\n}\n" }, { "path": "SAM2-Demo/Common/CGImage+Extension.swift", "content": "//\n// CGImage+Extension.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 8/20/24.\n//\n\nimport ImageIO\n\nextension CGImage {\n func resized(to size: CGSize) -> CGImage? {\n let width: Int = Int(size.width)\n let height: Int = Int(size.height)\n\n let bytesPerPixel = self.bitsPerPixel / 8\n let destBytesPerRow = width * bytesPerPixel\n\n guard let colorSpace = self.colorSpace else { return nil }\n guard let context = CGContext(data: nil,\n width: width,\n height: height,\n bitsPerComponent: self.bitsPerComponent,\n bytesPerRow: destBytesPerRow,\n space: colorSpace,\n bitmapInfo: self.bitmapInfo.rawValue)\n else { return nil }\n\n context.interpolationQuality = .high\n context.draw(self, in: CGRect(x: 0, y: 0, width: width, height: height))\n\n return context.makeImage()\n }\n}\n" }, { "path": "SAM2-Demo/Common/CGImage+RawBytes.swift", "content": "/*\n Copyright (c) 2017-2019 M.I. Hollemans\n\n Permission is hereby granted, free of charge, to any person obtaining a copy\n of this software and associated documentation files (the \"Software\"), to\n deal in the Software without restriction, including without limitation the\n rights to use, copy, modify, merge, publish, distribute, sublicense, and/or\n sell copies of the Software, and to permit persons to whom the Software is\n furnished to do so, subject to the following conditions:\n\n The above copyright notice and this permission notice shall be included in\n all copies or substantial portions of the Software.\n\n THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS\n IN THE SOFTWARE.\n*/\n\nimport CoreGraphics\n\nextension CGImage {\n /**\n Converts the image into an array of RGBA bytes.\n */\n @nonobjc public func toByteArrayRGBA() -> [UInt8] {\n var bytes = [UInt8](repeating: 0, count: width * height * 4)\n bytes.withUnsafeMutableBytes { ptr in\n if let colorSpace = colorSpace,\n let context = CGContext(\n data: ptr.baseAddress,\n width: width,\n height: height,\n bitsPerComponent: bitsPerComponent,\n bytesPerRow: bytesPerRow,\n space: colorSpace,\n bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) {\n let rect = CGRect(x: 0, y: 0, width: width, height: height)\n context.draw(self, in: rect)\n }\n }\n return bytes\n }\n\n /**\n Creates a new CGImage from an array of RGBA bytes.\n */\n @nonobjc public class func fromByteArrayRGBA(_ bytes: [UInt8],\n width: Int,\n height: Int) -> CGImage? {\n return fromByteArray(bytes, width: width, height: height,\n bytesPerRow: width * 4,\n colorSpace: CGColorSpaceCreateDeviceRGB(),\n alphaInfo: .premultipliedLast)\n }\n\n /**\n Creates a new CGImage from an array of grayscale bytes.\n */\n @nonobjc public class func fromByteArrayGray(_ bytes: [UInt8],\n width: Int,\n height: Int) -> CGImage? {\n return fromByteArray(bytes, width: width, height: height,\n bytesPerRow: width,\n colorSpace: CGColorSpaceCreateDeviceGray(),\n alphaInfo: .none)\n }\n\n @nonobjc class func fromByteArray(_ bytes: [UInt8],\n width: Int,\n height: Int,\n bytesPerRow: Int,\n colorSpace: CGColorSpace,\n alphaInfo: CGImageAlphaInfo) -> CGImage? {\n return bytes.withUnsafeBytes { ptr in\n let context = CGContext(data: UnsafeMutableRawPointer(mutating: ptr.baseAddress!),\n width: width,\n height: height,\n bitsPerComponent: 8,\n bytesPerRow: bytesPerRow,\n space: colorSpace,\n bitmapInfo: alphaInfo.rawValue)\n return context?.makeImage()\n }\n }\n}\n" }, { "path": "SAM2-Demo/Common/Color+Extension.swift", "content": "//\n// Color+Extension.swift\n// SAM2-Demo\n//\n// Created by Fleetwood on 01/10/2024.\n//\n\nimport SwiftUI\n\n#if canImport(UIKit)\nimport UIKit\n#elseif canImport(AppKit)\nimport AppKit\n#endif\n\nextension Color {\n #if canImport(UIKit)\n var asNative: UIColor { UIColor(self) }\n #elseif canImport(AppKit)\n var asNative: NSColor { NSColor(self) }\n #endif\n\n var rgba: (red: CGFloat, green: CGFloat, blue: CGFloat, alpha: CGFloat) {\n let color = asNative.usingColorSpace(.deviceRGB)!\n var t = (CGFloat(), CGFloat(), CGFloat(), CGFloat())\n color.getRed(&t.0, green: &t.1, blue: &t.2, alpha: &t.3)\n return t\n }\n}\n\nfunc colorDistance(_ color1: Color, _ color2: Color) -> Double {\n let rgb1 = color1.rgba;\n let rgb2 = color2.rgba;\n\n let rDiff = rgb1.red - rgb2.red\n let gDiff = rgb1.green - rgb2.green\n let bDiff = rgb1.blue - rgb2.blue\n\n return sqrt(rDiff*rDiff + gDiff*gDiff + bDiff*bDiff)\n}\n\n// Determine the Euclidean distance of all candidates from current set of colors.\n// Find the **maximum min-distance** from all current colors.\nfunc furthestColor(from existingColors: [Color], among candidateColors: [Color]) -> Color {\n var maxMinDistance: Double = 0\n var furthestColor: Color = SAMSegmentation.randomCandidateColor() ?? SAMSegmentation.defaultColor\n\n for candidate in candidateColors {\n let minDistance = existingColors.map { colorDistance(candidate, $0) }.min() ?? 0\n if minDistance > maxMinDistance {\n maxMinDistance = minDistance\n furthestColor = candidate\n }\n }\n\n return furthestColor\n}\n" }, { "path": "SAM2-Demo/Common/CoreImageExtensions.swift", "content": "import CoreImage\nimport CoreImage.CIFilterBuiltins\nimport ImageIO\nimport UniformTypeIdentifiers\n\nextension CIImage {\n /// Returns a resized image.\n func resized(to size: CGSize) -> CIImage {\n let outputScaleX = size.width / extent.width\n let outputScaleY = size.height / extent.height\n var outputImage = self.transformed(by: CGAffineTransform(scaleX: outputScaleX, y: outputScaleY))\n outputImage = outputImage.transformed(\n by: CGAffineTransform(translationX: -outputImage.extent.origin.x, y: -outputImage.extent.origin.y)\n )\n return outputImage\n }\n \n public func withAlpha(_ alpha: T) -> CIImage? {\n guard alpha != 1 else { return self }\n \n let filter = CIFilter.colorMatrix()\n filter.inputImage = self\n filter.aVector = CIVector(x: 0, y: 0, z: 0, w: CGFloat(alpha))\n\n return filter.outputImage\n }\n\n public func applyingThreshold(_ threshold: Float) -> CIImage? {\n let filter = CIFilter.colorThreshold()\n filter.inputImage = self\n filter.threshold = threshold\n return filter.outputImage\n }\n}\n\nextension CIContext {\n /// Renders an image to a new pixel buffer.\n func render(_ image: CIImage, pixelFormat: OSType) -> CVPixelBuffer? {\n var output: CVPixelBuffer!\n let status = CVPixelBufferCreate(\n kCFAllocatorDefault,\n Int(image.extent.width),\n Int(image.extent.height),\n pixelFormat,\n nil,\n &output\n )\n guard status == kCVReturnSuccess else {\n return nil\n }\n render(image, to: output, bounds: image.extent, colorSpace: nil)\n return output\n }\n\n /// Writes the image as a PNG.\n func writePNG(_ image: CIImage, to url: URL) {\n let outputCGImage = createCGImage(image, from: image.extent, format: .BGRA8, colorSpace: nil)!\n guard let destination = CGImageDestinationCreateWithURL(url as CFURL, UTType.png.identifier as CFString, 1, nil) else {\n fatalError(\"Failed to create an image destination.\")\n }\n CGImageDestinationAddImage(destination, outputCGImage, nil)\n CGImageDestinationFinalize(destination)\n }\n}\n" }, { "path": "SAM2-Demo/Common/DirectoryDocument.swift", "content": "//\n// DirectoryDocument.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/10/24.\n//\n\n\nimport SwiftUI\nimport UniformTypeIdentifiers\n\nstruct DirectoryDocument: FileDocument {\n static var readableContentTypes: [UTType] { [.folder] }\n\n init(initialContentType: UTType = .folder) {\n // Initialize if needed\n }\n\n init(configuration: ReadConfiguration) throws {\n // Initialize if needed\n }\n\n func fileWrapper(configuration: WriteConfiguration) throws -> FileWrapper {\n return FileWrapper(directoryWithFileWrappers: [:])\n }\n}\n" }, { "path": "SAM2-Demo/Common/MLMultiArray+Image.swift", "content": "/*\n Copyright (c) 2017-2020 M.I. Hollemans\n\n Permission is hereby granted, free of charge, to any person obtaining a copy\n of this software and associated documentation files (the \"Software\"), to\n deal in the Software without restriction, including without limitation the\n rights to use, copy, modify, merge, publish, distribute, sublicense, and/or\n sell copies of the Software, and to permit persons to whom the Software is\n furnished to do so, subject to the following conditions:\n\n The above copyright notice and this permission notice shall be included in\n all copies or substantial portions of the Software.\n\n THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING\n FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS\n IN THE SOFTWARE.\n*/\n\nimport Accelerate\nimport CoreML\n\npublic func clamp(_ x: T, min: T, max: T) -> T {\n if x < min { return min }\n if x > max { return max }\n return x\n}\n\npublic protocol MultiArrayType: Comparable {\n static var multiArrayDataType: MLMultiArrayDataType { get }\n static func +(lhs: Self, rhs: Self) -> Self\n static func -(lhs: Self, rhs: Self) -> Self\n static func *(lhs: Self, rhs: Self) -> Self\n static func /(lhs: Self, rhs: Self) -> Self\n init(_: Int)\n var toUInt8: UInt8 { get }\n}\n\nextension Double: MultiArrayType {\n public static var multiArrayDataType: MLMultiArrayDataType { return .double }\n public var toUInt8: UInt8 { return UInt8(self) }\n}\n\nextension Float: MultiArrayType {\n public static var multiArrayDataType: MLMultiArrayDataType { return .float32 }\n public var toUInt8: UInt8 { return UInt8(self) }\n}\n\nextension Int32: MultiArrayType {\n public static var multiArrayDataType: MLMultiArrayDataType { return .int32 }\n public var toUInt8: UInt8 { return UInt8(self) }\n}\n\nextension Float16: MultiArrayType {\n public static var multiArrayDataType: MLMultiArrayDataType { return .float16 }\n public var toUInt8: UInt8 { return UInt8(self) }\n}\n\nextension MLMultiArray {\n /**\n Converts the multi-array to a CGImage.\n\n The multi-array must have at least 2 dimensions for a grayscale image, or\n at least 3 dimensions for a color image.\n\n The default expected shape is (height, width) or (channels, height, width).\n However, you can change this using the `axes` parameter. For example, if\n the array shape is (1, height, width, channels), use `axes: (3, 1, 2)`.\n\n If `channel` is not nil, only converts that channel to a grayscale image.\n This lets you visualize individual channels from a multi-array with more\n than 4 channels.\n\n Otherwise, converts all channels. In this case, the number of channels in\n the multi-array must be 1 for grayscale, 3 for RGB, or 4 for RGBA.\n\n Use the `min` and `max` parameters to put the values from the array into\n the range [0, 255], if not already:\n\n - `min`: should be the smallest value in the data; this will be mapped to 0.\n - `max`: should be the largest value in the data; will be mapped to 255.\n\n For example, if the range of the data in the multi-array is [-1, 1], use\n `min: -1, max: 1`. If the range is already [0, 255], then use the defaults.\n */\n public func cgImage(min: Double = 0,\n max: Double = 255,\n channel: Int? = nil,\n axes: (Int, Int, Int)? = nil) -> CGImage? {\n switch self.dataType {\n case .double:\n return _image(min: min, max: max, channel: channel, axes: axes)\n case .float32:\n return _image(min: Float(min), max: Float(max), channel: channel, axes: axes)\n case .int32:\n return _image(min: Int32(min), max: Int32(max), channel: channel, axes: axes)\n case .float16:\n return _image(min: Float16(min), max: Float16(max), channel: channel, axes: axes)\n @unknown default:\n fatalError(\"Unsupported data type \\(dataType.rawValue)\")\n }\n }\n\n /**\n Helper function that allows us to use generics. The type of `min` and `max`\n is also the dataType of the MLMultiArray.\n */\n private func _image(min: T,\n max: T,\n channel: Int?,\n axes: (Int, Int, Int)?) -> CGImage? {\n if let (b, w, h, c) = toRawBytes(min: min, max: max, channel: channel, axes: axes) {\n if c == 1 {\n return CGImage.fromByteArrayGray(b, width: w, height: h)\n } else {\n return CGImage.fromByteArrayRGBA(b, width: w, height: h)\n }\n }\n return nil\n }\n\n /**\n Converts the multi-array into an array of RGBA or grayscale pixels.\n\n - Note: This is not particularly fast, but it is flexible. You can change\n the loops to convert the multi-array whichever way you please.\n\n - Note: The type of `min` and `max` must match the dataType of the\n MLMultiArray object.\n\n - Returns: tuple containing the RGBA bytes, the dimensions of the image,\n and the number of channels in the image (1, 3, or 4).\n */\n public func toRawBytes(min: T,\n max: T,\n channel: Int? = nil,\n axes: (Int, Int, Int)? = nil)\n -> (bytes: [UInt8], width: Int, height: Int, channels: Int)? {\n // MLMultiArray with unsupported shape?\n if shape.count < 2 {\n print(\"Cannot convert MLMultiArray of shape \\(shape) to image\")\n return nil\n }\n\n // Figure out which dimensions to use for the channels, height, and width.\n let channelAxis: Int\n let heightAxis: Int\n let widthAxis: Int\n if let axes = axes {\n channelAxis = axes.0\n heightAxis = axes.1\n widthAxis = axes.2\n guard channelAxis >= 0 && channelAxis < shape.count &&\n heightAxis >= 0 && heightAxis < shape.count &&\n widthAxis >= 0 && widthAxis < shape.count else {\n print(\"Invalid axes \\(axes) for shape \\(shape)\")\n return nil\n }\n } else if shape.count == 2 {\n // Expected shape for grayscale is (height, width)\n heightAxis = 0\n widthAxis = 1\n channelAxis = -1 // Never be used\n } else {\n // Expected shape for color is (channels, height, width)\n channelAxis = 0\n heightAxis = 1\n widthAxis = 2\n }\n\n let height = self.shape[heightAxis].intValue\n let width = self.shape[widthAxis].intValue\n let yStride = self.strides[heightAxis].intValue\n let xStride = self.strides[widthAxis].intValue\n\n let channels: Int\n let cStride: Int\n let bytesPerPixel: Int\n let channelOffset: Int\n\n // MLMultiArray with just two dimensions is always grayscale. (We ignore\n // the value of channelAxis here.)\n if shape.count == 2 {\n channels = 1\n cStride = 0\n bytesPerPixel = 1\n channelOffset = 0\n\n // MLMultiArray with more than two dimensions can be color or grayscale.\n } else {\n let channelDim = self.shape[channelAxis].intValue\n if let channel = channel {\n if channel < 0 || channel >= channelDim {\n print(\"Channel must be -1, or between 0 and \\(channelDim - 1)\")\n return nil\n }\n channels = 1\n bytesPerPixel = 1\n channelOffset = channel\n } else if channelDim == 1 {\n channels = 1\n bytesPerPixel = 1\n channelOffset = 0\n } else {\n if channelDim != 3 && channelDim != 4 {\n print(\"Expected channel dimension to have 1, 3, or 4 channels, got \\(channelDim)\")\n return nil\n }\n channels = channelDim\n bytesPerPixel = 4\n channelOffset = 0\n }\n cStride = self.strides[channelAxis].intValue\n }\n\n // Allocate storage for the RGBA or grayscale pixels. Set everything to\n // 255 so that alpha channel is filled in if only 3 channels.\n let count = height * width * bytesPerPixel\n var pixels = [UInt8](repeating: 255, count: count)\n\n // Grab the pointer to MLMultiArray's memory.\n var ptr = UnsafeMutablePointer(OpaquePointer(self.dataPointer))\n ptr = ptr.advanced(by: channelOffset * cStride)\n\n // Loop through all the pixels and all the channels and copy them over.\n for c in 0.. CGImage? {\n assert(features.dataType == .float32)\n assert(features.shape.count == 3)\n\n let ptr = UnsafeMutablePointer(OpaquePointer(features.dataPointer))\n\n let height = features.shape[1].intValue\n let width = features.shape[2].intValue\n let channelStride = features.strides[0].intValue\n let rowStride = features.strides[1].intValue\n let srcRowBytes = rowStride * MemoryLayout.stride\n\n var blueBuffer = vImage_Buffer(data: ptr,\n height: vImagePixelCount(height),\n width: vImagePixelCount(width),\n rowBytes: srcRowBytes)\n var greenBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride),\n height: vImagePixelCount(height),\n width: vImagePixelCount(width),\n rowBytes: srcRowBytes)\n var redBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride * 2),\n height: vImagePixelCount(height),\n width: vImagePixelCount(width),\n rowBytes: srcRowBytes)\n\n let destRowBytes = width * 4\n\n var error: vImage_Error = 0\n var pixels = [UInt8](repeating: 0, count: height * destRowBytes)\n\n pixels.withUnsafeMutableBufferPointer { ptr in\n var destBuffer = vImage_Buffer(data: ptr.baseAddress!,\n height: vImagePixelCount(height),\n width: vImagePixelCount(width),\n rowBytes: destRowBytes)\n\n error = vImageConvert_PlanarFToBGRX8888(&blueBuffer,\n &greenBuffer,\n &redBuffer,\n Pixel_8(255),\n &destBuffer,\n [max, max, max],\n [min, min, min],\n vImage_Flags(0))\n }\n\n if error == kvImageNoError {\n return CGImage.fromByteArrayRGBA(pixels, width: width, height: height)\n } else {\n return nil\n }\n}\n" }, { "path": "SAM2-Demo/Common/Models.swift", "content": "//\n// Models.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 8/19/24.\n//\n\nimport Foundation\nimport SwiftUI\n\nenum SAMCategoryType: Int {\n case background = 0\n case foreground = 1\n case boxOrigin = 2\n case boxEnd = 3\n\n var description: String {\n switch self {\n case .foreground:\n return \"Foreground\"\n case .background:\n return \"Background\"\n case .boxOrigin:\n return \"Box Origin\"\n case .boxEnd:\n return \"Box End\"\n }\n }\n}\n\nstruct SAMCategory: Hashable {\n let id: UUID = UUID()\n let type: SAMCategoryType\n let name: String\n let iconName: String\n let color: Color\n\n var typeDescription: String {\n type.description\n }\n\n static let foreground = SAMCategory(\n type: .foreground,\n name: \"Foreground\",\n iconName: \"square.on.square.dashed\",\n color: .pink\n )\n\n static let background = SAMCategory(\n type: .background,\n name: \"Background\",\n iconName: \"square.on.square.intersection.dashed\",\n color: .purple\n )\n\n static let boxOrigin = SAMCategory(\n type: .boxOrigin,\n name: \"Box Origin\",\n iconName: \"\",\n color: .white\n )\n\n static let boxEnd = SAMCategory(\n type: .boxEnd,\n name: \"Box End\",\n iconName: \"\",\n color: .white\n )\n}\n\n\nstruct SAMPoint: Hashable {\n let id = UUID()\n let coordinates: CGPoint\n let category: SAMCategory\n let dateAdded = Date()\n}\n\nstruct SAMBox: Hashable, Identifiable {\n let id = UUID()\n var startPoint: CGPoint\n var endPoint: CGPoint\n let category: SAMCategory\n let dateAdded = Date()\n var midpoint: CGPoint {\n return CGPoint(\n x: (startPoint.x + endPoint.x) / 2,\n y: (startPoint.y + endPoint.y) / 2\n )\n }\n}\n\nextension SAMBox {\n var points: [SAMPoint] {\n [SAMPoint(coordinates: startPoint, category: .boxOrigin), SAMPoint(coordinates: endPoint, category: .boxEnd)]\n }\n}\n\nstruct SAMSegmentation: Hashable, Identifiable {\n let id = UUID()\n var image: CIImage\n var tintColor: Color {\n didSet {\n updateTintedImage()\n }\n }\n var title: String = \"\"\n var firstAppearance: Int?\n var isHidden: Bool = false\n \n private var tintedImage: CIImage?\n\n static let defaultColor: Color = Color(.sRGB, red: 30/255, green: 144/255, blue: 1)\n static let candidateColors: [Color] = [\n defaultColor,\n Color.red,\n Color.green,\n Color.brown,\n Color.indigo,\n Color.cyan,\n Color.yellow,\n Color.purple,\n Color.orange,\n Color.teal,\n Color.indigo,\n Color.mint,\n Color.pink,\n ]\n\n init(image: CIImage, tintColor: Color = Color(.sRGB, red: 30/255, green: 144/255, blue: 1), title: String = \"\", firstAppearance: Int? = nil, isHidden: Bool = false) {\n self.image = image\n self.tintColor = tintColor\n self.title = title\n self.firstAppearance = firstAppearance\n self.isHidden = isHidden\n updateTintedImage()\n }\n\n private mutating func updateTintedImage() {\n let ciColor = CIColor(color: NSColor(tintColor))\n let monochromeFilter = CIFilter.colorMonochrome()\n monochromeFilter.inputImage = image\n monochromeFilter.color = ciColor!\n monochromeFilter.intensity = 1.0\n tintedImage = monochromeFilter.outputImage\n }\n \n static func randomCandidateColor() -> Color? {\n Self.candidateColors.randomElement()\n }\n\n var cgImage: CGImage {\n let context = CIContext()\n return context.createCGImage(tintedImage ?? image, from: (tintedImage ?? image).extent)!\n }\n}\n\nstruct SAMTool: Hashable {\n let id: UUID = UUID()\n let name: String\n let iconName: String\n}\n\n// Tools\nlet pointTool: SAMTool = SAMTool(name: \"Point\", iconName: \"hand.point.up.left\")\nlet boundingBoxTool: SAMTool = SAMTool(name: \"Bounding Box\", iconName: \"rectangle.dashed\")\n" }, { "path": "SAM2-Demo/Common/NSImage+Extension.swift", "content": "//\n// NSImage+Extension.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 8/20/24.\n//\n\nimport AppKit\nimport VideoToolbox\n\nextension NSImage {\n /**\n Converts the image to an ARGB `CVPixelBuffer`.\n */\n public func pixelBuffer() -> CVPixelBuffer? {\n return pixelBuffer(width: Int(size.width), height: Int(size.height))\n }\n\n /**\n Resizes the image to `width` x `height` and converts it to an ARGB\n `CVPixelBuffer`.\n */\n public func pixelBuffer(width: Int, height: Int) -> CVPixelBuffer? {\n return pixelBuffer(width: width, height: height,\n pixelFormatType: kCVPixelFormatType_32ARGB,\n colorSpace: CGColorSpaceCreateDeviceRGB(),\n alphaInfo: .noneSkipFirst)\n }\n\n /**\n Resizes the image to `width` x `height` and converts it to a `CVPixelBuffer`\n with the specified pixel format, color space, and alpha channel.\n */\n public func pixelBuffer(width: Int, height: Int,\n pixelFormatType: OSType,\n colorSpace: CGColorSpace,\n alphaInfo: CGImageAlphaInfo) -> CVPixelBuffer? {\n var maybePixelBuffer: CVPixelBuffer?\n let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,\n kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue]\n let status = CVPixelBufferCreate(kCFAllocatorDefault,\n width,\n height,\n pixelFormatType,\n attrs as CFDictionary,\n &maybePixelBuffer)\n\n guard status == kCVReturnSuccess, let pixelBuffer = maybePixelBuffer else {\n return nil\n }\n\n let flags = CVPixelBufferLockFlags(rawValue: 0)\n guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(pixelBuffer, flags) else {\n return nil\n }\n defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, flags) }\n\n guard let context = CGContext(data: CVPixelBufferGetBaseAddress(pixelBuffer),\n width: width,\n height: height,\n bitsPerComponent: 8,\n bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer),\n space: colorSpace,\n bitmapInfo: alphaInfo.rawValue)\n else {\n return nil\n }\n\n NSGraphicsContext.saveGraphicsState()\n let nscg = NSGraphicsContext(cgContext: context, flipped: true)\n NSGraphicsContext.current = nscg\n context.translateBy(x: 0, y: CGFloat(height))\n context.scaleBy(x: 1, y: -1)\n self.draw(in: CGRect(x: 0, y: 0, width: width, height: height))\n NSGraphicsContext.restoreGraphicsState()\n\n return pixelBuffer\n }\n}\n" }, { "path": "SAM2-Demo/Common/SAM2.swift", "content": "//\n// SAM2.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 8/20/24.\n//\n\nimport SwiftUI\nimport CoreML\nimport CoreImage\nimport CoreImage.CIFilterBuiltins\nimport Combine\nimport UniformTypeIdentifiers\n\n@MainActor\nclass SAM2: ObservableObject {\n \n @Published var imageEncodings: SAM2_1SmallImageEncoderFLOAT16Output?\n @Published var promptEncodings: SAM2_1SmallPromptEncoderFLOAT16Output?\n\n @Published private(set) var initializationTime: TimeInterval?\n @Published private(set) var initialized: Bool?\n\n private var imageEncoderModel: SAM2_1SmallImageEncoderFLOAT16?\n private var promptEncoderModel: SAM2_1SmallPromptEncoderFLOAT16?\n private var maskDecoderModel: SAM2_1SmallMaskDecoderFLOAT16?\n\n // TODO: examine model inputs instead\n var inputSize: CGSize { CGSize(width: 1024, height: 1024) }\n var width: CGFloat { inputSize.width }\n var height: CGFloat { inputSize.height }\n\n init() {\n Task {\n await loadModels()\n }\n }\n \n private func loadModels() async {\n let startTime = CFAbsoluteTimeGetCurrent()\n \n do {\n let configuration = MLModelConfiguration()\n configuration.computeUnits = .cpuAndGPU\n let (imageEncoder, promptEncoder, maskDecoder) = try await Task.detached(priority: .userInitiated) {\n let imageEncoder = try SAM2_1SmallImageEncoderFLOAT16(configuration: configuration)\n let promptEncoder = try SAM2_1SmallPromptEncoderFLOAT16(configuration: configuration)\n let maskDecoder = try SAM2_1SmallMaskDecoderFLOAT16(configuration: configuration)\n return (imageEncoder, promptEncoder, maskDecoder)\n }.value\n \n let endTime = CFAbsoluteTimeGetCurrent()\n self.initializationTime = endTime - startTime\n self.initialized = true\n\n self.imageEncoderModel = imageEncoder\n self.promptEncoderModel = promptEncoder\n self.maskDecoderModel = maskDecoder\n print(\"Initialized models in \\(String(format: \"%.4f\", self.initializationTime!)) seconds\")\n } catch {\n print(\"Failed to initialize models: \\(error)\")\n self.initializationTime = nil\n self.initialized = false\n }\n }\n\n // Convenience for use in the CLI\n private var modelLoading: AnyCancellable?\n func ensureModelsAreLoaded() async throws -> SAM2 {\n let _ = try await withCheckedThrowingContinuation { continuation in\n modelLoading = self.$initialized.sink { newValue in\n if let initialized = newValue {\n if initialized {\n continuation.resume(returning: self)\n } else {\n continuation.resume(throwing: SAM2Error.modelNotLoaded)\n }\n }\n }\n }\n return self\n }\n\n static func load() async throws -> SAM2 {\n try await SAM2().ensureModelsAreLoaded()\n }\n\n func getImageEncoding(from pixelBuffer: CVPixelBuffer) async throws {\n guard let model = imageEncoderModel else {\n throw SAM2Error.modelNotLoaded\n }\n \n let encoding = try model.prediction(image: pixelBuffer)\n self.imageEncodings = encoding\n }\n\n func getImageEncoding(from url: URL) async throws {\n guard let model = imageEncoderModel else {\n throw SAM2Error.modelNotLoaded\n }\n\n let inputs = try SAM2_1SmallImageEncoderFLOAT16Input(imageAt: url)\n let encoding = try await model.prediction(input: inputs)\n self.imageEncodings = encoding\n }\n\n func getPromptEncoding(from allPoints: [SAMPoint], with size: CGSize) async throws {\n guard let model = promptEncoderModel else {\n throw SAM2Error.modelNotLoaded\n }\n \n let transformedCoords = try transformCoords(allPoints.map { $0.coordinates }, normalize: false, origHW: size)\n\n // Create MLFeatureProvider with the required input format\n let pointsMultiArray = try MLMultiArray(shape: [1, NSNumber(value: allPoints.count), 2], dataType: .float32)\n let labelsMultiArray = try MLMultiArray(shape: [1, NSNumber(value: allPoints.count)], dataType: .int32)\n \n for (index, point) in transformedCoords.enumerated() {\n pointsMultiArray[[0, index, 0] as [NSNumber]] = NSNumber(value: Float(point.x))\n pointsMultiArray[[0, index, 1] as [NSNumber]] = NSNumber(value: Float(point.y))\n labelsMultiArray[[0, index] as [NSNumber]] = NSNumber(value: allPoints[index].category.type.rawValue)\n }\n \n let encoding = try model.prediction(points: pointsMultiArray, labels: labelsMultiArray)\n self.promptEncodings = encoding\n }\n\n func bestMask(for output: SAM2_1SmallMaskDecoderFLOAT16Output) -> MLMultiArray {\n if #available(macOS 15.0, *) {\n let scores = output.scoresShapedArray.scalars\n let argmax = scores.firstIndex(of: scores.max() ?? 0) ?? 0\n return MLMultiArray(output.low_res_masksShapedArray[0, argmax])\n } else {\n // Convert scores to float32 for compatibility with macOS < 15,\n // plus ugly loop copy (could do some memcpys)\n let scores = output.scores\n let floatScores = (0.. CIImage? {\n guard let model = maskDecoderModel else {\n throw SAM2Error.modelNotLoaded\n }\n \n if let image_embedding = self.imageEncodings?.image_embedding,\n let feats0 = self.imageEncodings?.feats_s0,\n let feats1 = self.imageEncodings?.feats_s1,\n let sparse_embedding = self.promptEncodings?.sparse_embeddings,\n let dense_embedding = self.promptEncodings?.dense_embeddings {\n let output = try model.prediction(image_embedding: image_embedding, sparse_embedding: sparse_embedding, dense_embedding: dense_embedding, feats_s0: feats0, feats_s1: feats1)\n\n // Extract best mask and ignore the others\n let lowFeatureMask = bestMask(for: output)\n\n // TODO: optimization\n // Preserve range for upsampling\n var minValue: Double = 9999\n var maxValue: Double = -9999\n for i in 0.. maxValue { maxValue = v }\n if v < minValue { minValue = v }\n }\n let threshold = -minValue / (maxValue - minValue)\n\n // Resize first, then threshold\n if let maskcgImage = lowFeatureMask.cgImage(min: minValue, max: maxValue) {\n let ciImage = CIImage(cgImage: maskcgImage, options: [.colorSpace: NSNull()])\n let resizedImage = try resizeImage(ciImage, to: original_size, applyingThreshold: Float(threshold))\n return resizedImage?.maskedToAlpha()?.samTinted()\n }\n }\n return nil\n }\n\n private func transformCoords(_ coords: [CGPoint], normalize: Bool = false, origHW: CGSize) throws -> [CGPoint] {\n guard normalize else {\n return coords.map { CGPoint(x: $0.x * width, y: $0.y * height) }\n }\n \n let w = origHW.width\n let h = origHW.height\n \n return coords.map { coord in\n let normalizedX = coord.x / w\n let normalizedY = coord.y / h\n return CGPoint(x: normalizedX * width, y: normalizedY * height)\n }\n }\n \n private func resizeImage(_ image: CIImage, to size: CGSize, applyingThreshold threshold: Float = 1) throws -> CIImage? {\n let scale = CGAffineTransform(scaleX: size.width / image.extent.width,\n y: size.height / image.extent.height)\n return image.transformed(by: scale).applyingThreshold(threshold)\n }\n}\n\nextension CIImage {\n /// This is only appropriate for grayscale mask images (our case). CIColorMatrix can be used more generally.\n func maskedToAlpha() -> CIImage? {\n let filter = CIFilter.maskToAlpha()\n filter.inputImage = self\n return filter.outputImage\n }\n\n func samTinted() -> CIImage? {\n let filter = CIFilter.colorMatrix()\n filter.rVector = CIVector(x: 30/255, y: 0, z: 0, w: 1)\n filter.gVector = CIVector(x: 0, y: 144/255, z: 0, w: 1)\n filter.bVector = CIVector(x: 0, y: 0, z: 1, w: 1)\n filter.biasVector = CIVector(x: -1, y: -1, z: -1, w: 0)\n filter.inputImage = self\n return filter.outputImage?.cropped(to: self.extent)\n }\n}\n\nenum SAM2Error: Error {\n case modelNotLoaded\n case pixelBufferCreationFailed\n case imageResizingFailed\n}\n\n@discardableResult func writeCGImage(_ image: CGImage, to destinationURL: URL) -> Bool {\n guard let destination = CGImageDestinationCreateWithURL(destinationURL as CFURL, UTType.png.identifier as CFString, 1, nil) else { return false }\n CGImageDestinationAddImage(destination, image, nil)\n return CGImageDestinationFinalize(destination)\n}\n" }, { "path": "SAM2-Demo/ContentView.swift", "content": "import SwiftUI\nimport PhotosUI\nimport UniformTypeIdentifiers\nimport CoreML\n\nimport os\n\n// TODO: Add reset, bounding box, and eraser\n\nlet logger = Logger(\n subsystem:\n \"com.cyrilzakka.SAM2-Demo.ContentView\",\n category: \"ContentView\")\n\n\nstruct PointsOverlay: View {\n @Binding var selectedPoints: [SAMPoint]\n @Binding var selectedTool: SAMTool?\n let imageSize: CGSize\n \n var body: some View {\n ForEach(selectedPoints, id: \\.self) { point in\n Circle()\n .frame(width: 10, height: 10)\n .foregroundStyle(point.category.color)\n .position(point.coordinates.toSize(imageSize))\n \n }\n }\n}\n\nstruct BoundingBoxesOverlay: View {\n let boundingBoxes: [SAMBox]\n let currentBox: SAMBox?\n let imageSize: CGSize\n \n var body: some View {\n ForEach(boundingBoxes) { box in\n BoundingBoxPath(box: box, imageSize: imageSize)\n }\n if let currentBox = currentBox {\n BoundingBoxPath(box: currentBox, imageSize: imageSize)\n }\n }\n}\n\nstruct BoundingBoxPath: View {\n let box: SAMBox\n let imageSize: CGSize\n \n var body: some View {\n Path { path in\n path.move(to: box.startPoint.toSize(imageSize))\n path.addLine(to: CGPoint(x: box.endPoint.x, y: box.startPoint.y).toSize(imageSize))\n path.addLine(to: box.endPoint.toSize(imageSize))\n path.addLine(to: CGPoint(x: box.startPoint.x, y: box.endPoint.y).toSize(imageSize))\n path.closeSubpath()\n }\n .stroke(\n box.category.color,\n style: StrokeStyle(lineWidth: 2, dash: [5, 5])\n )\n }\n}\n\nstruct SegmentationOverlay: View {\n \n @Binding var segmentationImage: SAMSegmentation\n let imageSize: CGSize\n \n @State var counter: Int = 0\n var origin: CGPoint = .zero\n var shouldAnimate: Bool = false\n \n var body: some View {\n let nsImage = NSImage(cgImage: segmentationImage.cgImage, size: imageSize)\n Image(nsImage: nsImage)\n .resizable()\n .scaledToFit()\n .allowsHitTesting(false)\n .frame(width: imageSize.width, height: imageSize.height)\n .opacity(segmentationImage.isHidden ? 0:0.6)\n .modifier(RippleEffect(at: CGPoint(x: segmentationImage.cgImage.width/2, y: segmentationImage.cgImage.height/2), trigger: counter))\n .onAppear {\n if shouldAnimate {\n counter += 1\n }\n }\n }\n}\n\nstruct ContentView: View {\n \n // ML Models\n @StateObject private var sam2 = SAM2()\n @State private var currentSegmentation: SAMSegmentation?\n @State private var segmentationImages: [SAMSegmentation] = []\n @State private var imageSize: CGSize = .zero\n \n // File importer\n @State private var imageURL: URL?\n @State private var isImportingFromFiles: Bool = false\n @State private var displayImage: NSImage?\n \n // Mask exporter\n @State private var exportURL: URL?\n @State private var exportMaskToPNG: Bool = false\n @State private var showInspector: Bool = true\n @State private var selectedSegmentations = Set()\n \n // Photos Picker\n @State private var isImportingFromPhotos: Bool = false\n @State private var selectedItem: PhotosPickerItem?\n \n @State private var error: Error?\n \n // ML Model Properties\n var tools: [SAMTool] = [pointTool, boundingBoxTool]\n var categories: [SAMCategory] = [.foreground, .background]\n \n @State private var selectedTool: SAMTool?\n @State private var selectedCategory: SAMCategory?\n @State private var selectedPoints: [SAMPoint] = []\n @State private var boundingBoxes: [SAMBox] = []\n @State private var currentBox: SAMBox?\n @State private var originalSize: NSSize?\n @State private var currentScale: CGFloat = 1.0\n @State private var visibleRect: CGRect = .zero\n \n var body: some View {\n \n NavigationSplitView(sidebar: {\n VStack {\n LayerListView(segmentationImages: $segmentationImages, selectedSegmentations: $selectedSegmentations, currentSegmentation: $currentSegmentation)\n Spacer()\n Button(action: {\n if let currentSegmentation = self.currentSegmentation {\n self.segmentationImages.append(currentSegmentation)\n\n self.reset()\n }\n }, label: {\n Text(\"New Mask\")\n }).padding()\n }\n }, detail: {\n ZStack {\n ZoomableScrollView(visibleRect: $visibleRect) {\n if let image = displayImage {\n ImageView(image: image, currentScale: $currentScale, selectedTool: $selectedTool, selectedCategory: $selectedCategory, selectedPoints: $selectedPoints, boundingBoxes: $boundingBoxes, currentBox: $currentBox, segmentationImages: $segmentationImages, currentSegmentation: $currentSegmentation, imageSize: $imageSize, originalSize: $originalSize, sam2: sam2)\n } else {\n ContentUnavailableView(\"No Image Loaded\", systemImage: \"photo.fill.on.rectangle.fill\", description: Text(\"Please import a photo to get started.\"))\n }\n }\n VStack(spacing: 0) {\n SubToolbar(selectedPoints: $selectedPoints, boundingBoxes: $boundingBoxes, segmentationImages: $segmentationImages, currentSegmentation: $currentSegmentation)\n Spacer()\n }\n \n }\n })\n .inspector(isPresented: $showInspector, content: {\n if selectedSegmentations.isEmpty {\n ContentUnavailableView(label: {\n Label(title: {\n Text(\"No Mask Selected\")\n .font(.subheadline)\n }, icon: {})\n \n })\n .inspectorColumnWidth(min: 200, ideal: 200, max: 200)\n } else {\n MaskEditor(exportMaskToPNG: $exportMaskToPNG, segmentationImages: $segmentationImages, selectedSegmentations: $selectedSegmentations, currentSegmentation: $currentSegmentation)\n .inspectorColumnWidth(min: 200, ideal: 200, max: 200)\n .toolbar {\n Spacer()\n Button {\n showInspector.toggle()\n } label: {\n Label(\"Toggle Inspector\", systemImage: \"sidebar.trailing\")\n }\n }\n \n }\n \n })\n .toolbar {\n // Tools\n ToolbarItemGroup(placement: .principal) {\n Picker(selection: $selectedTool, content: {\n ForEach(tools, id: \\.self) { tool in\n Label(tool.name, systemImage: tool.iconName)\n .tag(tool)\n .labelStyle(.titleAndIcon)\n }\n }, label: {\n Label(\"Tools\", systemImage: \"pencil.and.ruler\")\n })\n .pickerStyle(.menu)\n \n Picker(selection: $selectedCategory, content: {\n ForEach(categories, id: \\.self) { cat in\n Label(cat.name, systemImage: cat.iconName)\n .tag(cat)\n .labelStyle(.titleAndIcon)\n }\n }, label: {\n Label(\"Tools\", systemImage: \"pencil.and.ruler\")\n })\n .pickerStyle(.menu)\n \n }\n \n // Import\n ToolbarItemGroup {\n Menu {\n Button(action: {\n isImportingFromPhotos = true\n }, label: {\n Label(\"From Photos\", systemImage: \"photo.on.rectangle.angled.fill\")\n })\n \n Button(action: {\n isImportingFromFiles = true\n }, label: {\n Label(\"From Files\", systemImage: \"folder.fill\")\n })\n } label: {\n Label(\"Import\", systemImage: \"photo.badge.plus\")\n }\n }\n }\n \n .onAppear {\n if selectedTool == nil {\n selectedTool = tools[0]\n }\n if selectedCategory == nil {\n selectedCategory = categories.first\n }\n \n }\n \n // MARK: - Image encoding\n .onChange(of: displayImage) {\n segmentationImages = []\n self.reset()\n Task {\n if let displayImage, let pixelBuffer = displayImage.pixelBuffer(width: 1024, height: 1024) {\n originalSize = displayImage.size\n do {\n try await sam2.getImageEncoding(from: pixelBuffer)\n } catch {\n self.error = error\n }\n }\n }\n }\n \n \n // MARK: - Photos Importer\n .photosPicker(isPresented: $isImportingFromPhotos, selection: $selectedItem, matching: .any(of: [.images, .screenshots, .livePhotos]))\n .onChange(of: selectedItem) {\n Task {\n if let loadedData = try? await\n selectedItem?.loadTransferable(type: Data.self) {\n DispatchQueue.main.async {\n selectedPoints.removeAll()\n displayImage = NSImage(data: loadedData)\n }\n } else {\n logger.error(\"Error loading image from Photos.\")\n }\n }\n }\n \n // MARK: - File Importer\n .fileImporter(isPresented: $isImportingFromFiles,\n allowedContentTypes: [.image]) { result in\n switch result {\n case .success(let file):\n self.selectedItem = nil\n self.selectedPoints.removeAll()\n self.imageURL = file\n loadImage(from: file)\n case .failure(let error):\n logger.error(\"File import error: \\(error.localizedDescription)\")\n self.error = error\n }\n }\n \n // MARK: - File exporter\n .fileExporter(\n isPresented: $exportMaskToPNG,\n document: DirectoryDocument(initialContentType: .folder),\n contentType: .folder,\n defaultFilename: \"Segmentations\"\n ) { result in\n if case .success(let url) = result {\n exportURL = url\n var selectedToExport = segmentationImages.filter { segmentation in\n selectedSegmentations.contains(segmentation.id)\n }\n if let currentSegmentation {\n selectedToExport.append(currentSegmentation)\n }\n exportSegmentations(selectedToExport, to: url)\n }\n }\n }\n \n // MARK: - Private Methods\n private func loadImage(from url: URL) {\n guard url.startAccessingSecurityScopedResource() else {\n logger.error(\"Failed to access the file. Security-scoped resource access denied.\")\n return\n }\n \n defer { url.stopAccessingSecurityScopedResource() }\n \n do {\n let imageData = try Data(contentsOf: url)\n if let image = NSImage(data: imageData) {\n DispatchQueue.main.async {\n self.displayImage = image\n }\n } else {\n logger.error(\"Failed to create NSImage from file data\")\n }\n } catch {\n logger.error(\"Error loading image data: \\(error.localizedDescription)\")\n self.error = error\n }\n }\n \n func exportSegmentations(_ segmentations: [SAMSegmentation], to directory: URL) {\n let fileManager = FileManager.default\n \n do {\n try fileManager.createDirectory(at: directory, withIntermediateDirectories: true, attributes: nil)\n \n for (index, segmentation) in segmentations.enumerated() {\n let filename = \"segmentation_\\(index + 1).png\"\n let fileURL = directory.appendingPathComponent(filename)\n \n if let destination = CGImageDestinationCreateWithURL(fileURL as CFURL, UTType.png.identifier as CFString, 1, nil) {\n CGImageDestinationAddImage(destination, segmentation.cgImage, nil)\n if CGImageDestinationFinalize(destination) {\n print(\"Saved segmentation \\(index + 1) to \\(fileURL.path)\")\n } else {\n print(\"Failed to save segmentation \\(index + 1)\")\n }\n }\n }\n } catch {\n print(\"Error creating directory: \\(error.localizedDescription)\")\n }\n }\n \n private func reset() {\n selectedPoints = []\n boundingBoxes = []\n currentBox = nil\n currentSegmentation = nil\n }\n \n \n}\n\nstruct SizePreferenceKey: PreferenceKey {\n static var defaultValue: CGSize = .zero\n static func reduce(value: inout CGSize, nextValue: () -> CGSize) {\n value = nextValue()\n }\n}\n\n#Preview {\n ContentView()\n}\n" }, { "path": "SAM2-Demo/Preview Content/Preview Assets.xcassets/Contents.json", "content": "{\n \"info\" : {\n \"author\" : \"xcode\",\n \"version\" : 1\n }\n}\n" }, { "path": "SAM2-Demo/Ripple/Ripple.metal", "content": "// Ripple.metal\n\n/*\nSee the LICENSE at the end of the article for this sample’s licensing information.\n\nAbstract:\nA shader that applies a ripple effect to a view when using it as a SwiftUI layer\n effect.\n*/\n\n#include \n#include \nusing namespace metal;\n\n[[ stitchable ]]\nhalf4 Ripple(\n float2 position,\n SwiftUI::Layer layer,\n float2 origin,\n float time,\n float amplitude,\n float frequency,\n float decay,\n float speed\n) {\n // The distance of the current pixel position from `origin`.\n float distance = length(position - origin);\n // The amount of time it takes for the ripple to arrive at the current pixel position.\n float delay = distance / speed;\n\n // Adjust for delay, clamp to 0.\n time -= delay;\n time = max(0.0, time);\n\n // The ripple is a sine wave that Metal scales by an exponential decay\n // function.\n float rippleAmount = amplitude * sin(frequency * time) * exp(-decay * time);\n\n // A vector of length `amplitude` that points away from position.\n float2 n = normalize(position - origin);\n\n // Scale `n` by the ripple amount at the current pixel position and add it\n // to the current pixel position.\n //\n // This new position moves toward or away from `origin` based on the\n // sign and magnitude of `rippleAmount`.\n float2 newPosition = position + rippleAmount * n;\n\n // Sample the layer at the new position.\n half4 color = layer.sample(newPosition);\n\n // Lighten or darken the color based on the ripple amount and its alpha\n // component.\n color.rgb += 0.3 * (rippleAmount / amplitude) * color.a;\n\n return color;\n}\n" }, { "path": "SAM2-Demo/Ripple/RippleModifier.swift", "content": "//\n// RippleModifier.swift\n// HuggingChat-Mac\n//\n// Created by Cyril Zakka on 8/28/24.\n//\n\nimport SwiftUI\n\n/* See the LICENSE at the end of the article for this sample's licensing information. */\n\n/// A modifier that applies a ripple effect to its content.\nstruct RippleModifier: ViewModifier {\n var origin: CGPoint\n\n var elapsedTime: TimeInterval\n\n var duration: TimeInterval\n\n var amplitude: Double\n var frequency: Double\n var decay: Double\n var speed: Double\n\n func body(content: Content) -> some View {\n let shader = ShaderLibrary.Ripple(\n .float2(origin),\n .float(elapsedTime),\n\n // Parameters\n .float(amplitude),\n .float(frequency),\n .float(decay),\n .float(speed)\n )\n\n let maxSampleOffset = maxSampleOffset\n let elapsedTime = elapsedTime\n let duration = duration\n\n content.visualEffect { view, _ in\n view.layerEffect(\n shader,\n maxSampleOffset: maxSampleOffset,\n isEnabled: 0 < elapsedTime && elapsedTime < duration\n )\n }\n }\n\n var maxSampleOffset: CGSize {\n CGSize(width: amplitude, height: amplitude)\n }\n}\n" }, { "path": "SAM2-Demo/Ripple/RippleViewModifier.swift", "content": "//\n// RippleViewModifier.swift\n// HuggingChat-Mac\n//\n// Created by Cyril Zakka on 8/28/24.\n//\n\nimport SwiftUI\n\nstruct RippleEffect: ViewModifier {\n \n var origin: CGPoint\n var trigger: T\n var amplitude: Double\n var frequency: Double\n var decay: Double\n var speed: Double\n\n init(at origin: CGPoint, trigger: T, amplitude: Double = 12, frequency: Double = 15, decay: Double = 8, speed: Double = 1200) {\n self.origin = origin\n self.trigger = trigger\n self.amplitude = amplitude\n self.frequency = frequency\n self.decay = decay\n self.speed = speed\n }\n\n func body(content: Content) -> some View {\n let origin = origin\n let duration = duration\n let amplitude = amplitude\n let frequency = frequency\n let decay = decay\n let speed = speed\n\n content.keyframeAnimator(\n initialValue: 0,\n trigger: trigger\n ) { view, elapsedTime in\n view.modifier(RippleModifier(\n origin: origin,\n elapsedTime: elapsedTime,\n duration: duration,\n amplitude: amplitude,\n frequency: frequency,\n decay: decay,\n speed: speed\n ))\n } keyframes: { _ in\n MoveKeyframe(0)\n LinearKeyframe(duration, duration: duration)\n }\n }\n\n var duration: TimeInterval { 3 }\n}\n" }, { "path": "SAM2-Demo/SAM2_1SmallImageEncoderFLOAT16.mlpackage/Manifest.json", "content": "{\n \"fileFormatVersion\": \"1.0.0\",\n \"itemInfoEntries\": {\n \"4C20C7AA-F42B-4CCD-84C3-73C031A91D48\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Weights\",\n \"name\": \"weights\",\n \"path\": \"com.apple.CoreML/weights\"\n },\n \"DDCB1D63-C7BD-4A13-8EB5-D7151371105B\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Specification\",\n \"name\": \"model.mlmodel\",\n \"path\": \"com.apple.CoreML/model.mlmodel\"\n }\n },\n \"rootModelIdentifier\": \"DDCB1D63-C7BD-4A13-8EB5-D7151371105B\"\n}\n" }, { "path": "SAM2-Demo/SAM2_1SmallMaskDecoderFLOAT16.mlpackage/Manifest.json", "content": "{\n \"fileFormatVersion\": \"1.0.0\",\n \"itemInfoEntries\": {\n \"6FA6762D-69A1-4A0B-AB0D-512638FD7ECF\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Specification\",\n \"name\": \"model.mlmodel\",\n \"path\": \"com.apple.CoreML/model.mlmodel\"\n },\n \"DB82D069-C4C9-41FB-A178-262063485D28\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Weights\",\n \"name\": \"weights\",\n \"path\": \"com.apple.CoreML/weights\"\n }\n },\n \"rootModelIdentifier\": \"6FA6762D-69A1-4A0B-AB0D-512638FD7ECF\"\n}\n" }, { "path": "SAM2-Demo/SAM2_1SmallPromptEncoderFLOAT16.mlpackage/Manifest.json", "content": "{\n \"fileFormatVersion\": \"1.0.0\",\n \"itemInfoEntries\": {\n \"BE0329D0-1E5D-4FF9-8ECE-350FC8DE699D\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Weights\",\n \"name\": \"weights\",\n \"path\": \"com.apple.CoreML/weights\"\n },\n \"C1F60EF7-4F31-4243-8BE5-C107CB23EADF\": {\n \"author\": \"com.apple.CoreML\",\n \"description\": \"CoreML Model Specification\",\n \"name\": \"model.mlmodel\",\n \"path\": \"com.apple.CoreML/model.mlmodel\"\n }\n },\n \"rootModelIdentifier\": \"C1F60EF7-4F31-4243-8BE5-C107CB23EADF\"\n}\n" }, { "path": "SAM2-Demo/SAM2_Demo.entitlements", "content": "\n\n\n\n\tcom.apple.security.app-sandbox\n\t\n\tcom.apple.security.files.user-selected.read-write\n\t\n\n\n" }, { "path": "SAM2-Demo/SAM2_DemoApp.swift", "content": "//\n// SAM2_DemoApp.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 8/19/24.\n//\n\nimport SwiftUI\n\n@main\nstruct SAM2_DemoApp: App {\n var body: some Scene {\n WindowGroup {\n ContentView()\n }\n .windowToolbarStyle(UnifiedWindowToolbarStyle(showsTitle: false))\n }\n}\n" }, { "path": "SAM2-Demo/Views/AnnotationListView.swift", "content": "//\n// AnnotationListView.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/8/24.\n//\n\nimport SwiftUI\n\nstruct AnnotationListView: View {\n \n @Binding var segmentation: SAMSegmentation\n @State var showHideIcon: Bool = false\n \n var body: some View {\n HStack {\n Image(nsImage: NSImage(cgImage: segmentation.cgImage, size: NSSize(width: 25, height: 25)))\n .background(.quinary)\n .mask(RoundedRectangle(cornerRadius: 5))\n \n VStack(alignment: .leading) {\n Text(segmentation.title)\n .font(.headline)\n .foregroundStyle(segmentation.isHidden ? .tertiary:.primary)\n// Text(segmentation.firstAppearance)\n// .font(.subheadline)\n// .foregroundStyle(.secondary)\n }\n Spacer()\n Button(\"\", systemImage: segmentation.isHidden ? \"eye.slash.fill\" :\"eye.fill\", action: {\n segmentation.isHidden.toggle()\n })\n .opacity(segmentation.isHidden ? 1 : (showHideIcon ? 1:0))\n .buttonStyle(.borderless)\n .foregroundStyle(.secondary)\n }\n .onHover { state in\n showHideIcon = state\n }\n }\n}\n" }, { "path": "SAM2-Demo/Views/ImageView.swift", "content": "//\n// ImageView.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/8/24.\n//\n\nimport SwiftUI\n\nstruct ImageView: View {\n let image: NSImage\n @Binding var currentScale: CGFloat\n @Binding var selectedTool: SAMTool?\n @Binding var selectedCategory: SAMCategory?\n @Binding var selectedPoints: [SAMPoint]\n @Binding var boundingBoxes: [SAMBox]\n @Binding var currentBox: SAMBox?\n @Binding var segmentationImages: [SAMSegmentation]\n @Binding var currentSegmentation: SAMSegmentation?\n @Binding var imageSize: CGSize\n @Binding var originalSize: NSSize?\n \n @State var animationPoint: CGPoint = .zero\n @ObservedObject var sam2: SAM2\n @State private var error: Error?\n \n var pointSequence: [SAMPoint] {\n boundingBoxes.flatMap { $0.points } + selectedPoints\n }\n\n var body: some View {\n \n Image(nsImage: image)\n .resizable()\n .aspectRatio(contentMode: .fit)\n .scaleEffect(currentScale)\n .onTapGesture(coordinateSpace: .local) { handleTap(at: $0) }\n .gesture(boundingBoxGesture)\n .onHover { changeCursorAppearance(is: $0) }\n .background(GeometryReader { geometry in\n Color.clear.preference(key: SizePreferenceKey.self, value: geometry.size)\n })\n .onPreferenceChange(SizePreferenceKey.self) { imageSize = $0 }\n .onChange(of: selectedPoints.count, {\n if !selectedPoints.isEmpty {\n performForwardPass()\n }\n })\n .onChange(of: boundingBoxes.count, {\n if !boundingBoxes.isEmpty {\n performForwardPass()\n }\n })\n .overlay {\n PointsOverlay(selectedPoints: $selectedPoints, selectedTool: $selectedTool, imageSize: imageSize)\n BoundingBoxesOverlay(boundingBoxes: boundingBoxes, currentBox: currentBox, imageSize: imageSize)\n\n if !segmentationImages.isEmpty {\n ForEach(Array(segmentationImages.enumerated()), id: \\.element.id) { index, segmentation in\n SegmentationOverlay(segmentationImage: $segmentationImages[index], imageSize: imageSize, shouldAnimate: false)\n .zIndex(Double (segmentationImages.count - index))\n }\n }\n \n if let currentSegmentation = currentSegmentation {\n SegmentationOverlay(segmentationImage: .constant(currentSegmentation), imageSize: imageSize, origin: animationPoint, shouldAnimate: true)\n .zIndex(Double(segmentationImages.count + 1))\n }\n }\n \n }\n \n private func changeCursorAppearance(is inside: Bool) {\n if inside {\n if selectedTool == pointTool {\n NSCursor.pointingHand.push()\n } else if selectedTool == boundingBoxTool {\n NSCursor.crosshair.push()\n }\n } else {\n NSCursor.pop()\n }\n }\n \n private var boundingBoxGesture: some Gesture {\n DragGesture(minimumDistance: 0)\n .onChanged { value in\n guard selectedTool == boundingBoxTool else { return }\n \n if currentBox == nil {\n currentBox = SAMBox(startPoint: value.startLocation.fromSize(imageSize), endPoint: value.location.fromSize(imageSize), category: selectedCategory!)\n } else {\n currentBox?.endPoint = value.location.fromSize(imageSize)\n }\n }\n .onEnded { value in\n guard selectedTool == boundingBoxTool else { return }\n \n if let box = currentBox {\n boundingBoxes.append(box)\n animationPoint = box.midpoint.toSize(imageSize)\n currentBox = nil\n }\n }\n }\n \n private func handleTap(at location: CGPoint) {\n if selectedTool == pointTool {\n placePoint(at: location)\n animationPoint = location\n }\n }\n \n private func placePoint(at coordinates: CGPoint) {\n let samPoint = SAMPoint(coordinates: coordinates.fromSize(imageSize), category: selectedCategory!)\n self.selectedPoints.append(samPoint)\n }\n \n private func performForwardPass() {\n Task {\n do {\n try await sam2.getPromptEncoding(from: pointSequence, with: imageSize)\n if let mask = try await sam2.getMask(for: originalSize ?? .zero) {\n DispatchQueue.main.async {\n let colorSet = self.segmentationImages.map { $0.tintColor };\n let furthestColor = furthestColor(from: colorSet, among: SAMSegmentation.candidateColors)\n let segmentationNumber = segmentationImages.count\n let segmentationOverlay = SAMSegmentation(image: mask, tintColor: furthestColor, title: \"Untitled \\(segmentationNumber + 1)\")\n self.currentSegmentation = segmentationOverlay\n }\n }\n } catch {\n self.error = error\n }\n }\n }\n}\n\n#Preview {\n ContentView()\n}\n\nextension CGPoint {\n func fromSize(_ size: CGSize) -> CGPoint {\n CGPoint(x: x / size.width, y: y / size.height)\n }\n\n func toSize(_ size: CGSize) -> CGPoint {\n CGPoint(x: x * size.width, y: y * size.height)\n }\n}\n" }, { "path": "SAM2-Demo/Views/LayerListView.swift", "content": "//\n// LayerListView.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/8/24.\n//\n\nimport SwiftUI\n\nstruct LayerListView: View {\n \n @Binding var segmentationImages: [SAMSegmentation]\n @Binding var selectedSegmentations: Set\n @Binding var currentSegmentation: SAMSegmentation?\n \n var body: some View {\n List(selection: $selectedSegmentations) {\n Section(\"Annotations List\") {\n ForEach(Array(segmentationImages.enumerated()), id: \\.element.id) { index, segmentation in\n AnnotationListView(segmentation: $segmentationImages[index])\n .padding(.horizontal, 5)\n .contextMenu {\n Button(role: .destructive) {\n if let index = segmentationImages.firstIndex(where: { $0.id == segmentation.id }) {\n segmentationImages.remove(at: index)\n }\n } label: {\n Label(\"Delete\", systemImage: \"trash.fill\")\n }\n }\n }\n .onDelete(perform: delete)\n .onMove(perform: move)\n\n if let currentSegmentation = currentSegmentation {\n AnnotationListView(segmentation: .constant(currentSegmentation))\n .tag(currentSegmentation.id)\n }\n \n \n }\n }\n .listStyle(.sidebar)\n }\n \n func delete(at offsets: IndexSet) {\n segmentationImages.remove(atOffsets: offsets)\n }\n \n func move(from source: IndexSet, to destination: Int) {\n segmentationImages.move(fromOffsets: source, toOffset: destination)\n }\n}\n\n#Preview {\n ContentView()\n}\n" }, { "path": "SAM2-Demo/Views/MaskEditor.swift", "content": "//\n// MaskEditor.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/10/24.\n//\n\nimport SwiftUI\n\nstruct MaskEditor: View {\n \n @Binding var exportMaskToPNG: Bool\n @Binding var segmentationImages: [SAMSegmentation]\n @Binding var selectedSegmentations: Set\n @Binding var currentSegmentation: SAMSegmentation?\n \n @State private var bgColor =\n Color(.sRGB, red: 30/255, green: 144/255, blue: 1)\n \n var body: some View {\n Form {\n Section {\n ColorPicker(\"Color\", selection: $bgColor)\n .onChange(of: bgColor) { oldColor, newColor in\n updateSelectedSegmentationsColor(newColor)\n }\n \n Button(\"Export Selected...\", action: {\n exportMaskToPNG = true\n })\n .frame(maxWidth: .infinity, maxHeight: .infinity)\n }\n \n }\n .frame(maxWidth: .infinity, maxHeight: .infinity)\n .onChange(of: selectedSegmentations) { oldValue, newValue in\n bgColor = getColorOfFirstSelectedSegmentation()\n }\n .onAppear {\n bgColor = getColorOfFirstSelectedSegmentation()\n }\n \n }\n \n private func updateSelectedSegmentationsColor(_ newColor: Color) {\n for id in selectedSegmentations {\n for index in segmentationImages.indices where segmentationImages[index].id == id {\n segmentationImages[index].tintColor = newColor\n }\n if currentSegmentation?.id == id {\n currentSegmentation?.tintColor = newColor\n }\n }\n }\n \n private func getColorOfFirstSelectedSegmentation() -> Color {\n if let firstSelectedId = selectedSegmentations.first {\n if let firstSelectedSegmentation = segmentationImages.first(where: { $0.id == firstSelectedId }) {\n return firstSelectedSegmentation.tintColor\n } else {\n if let currentSegmentation {\n return currentSegmentation.tintColor\n }\n }\n \n }\n return bgColor // Return default color if no segmentation is selected\n }\n}\n\n#Preview {\n ContentView()\n}\n" }, { "path": "SAM2-Demo/Views/SubtoolbarView.swift", "content": "//\n// SubtoolbarView.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/8/24.\n//\n\nimport SwiftUI\n\nstruct SubToolbar: View {\n @Binding var selectedPoints: [SAMPoint]\n @Binding var boundingBoxes: [SAMBox]\n @Binding var segmentationImages: [SAMSegmentation]\n @Binding var currentSegmentation: SAMSegmentation?\n\n var body: some View {\n if selectedPoints.count > 0 || boundingBoxes.count > 0 {\n ZStack {\n Rectangle()\n .fill(.regularMaterial)\n .frame(height: 30)\n \n HStack {\n Spacer()\n Button(\"Undo\", action: undo)\n .padding(.trailing, 5)\n .disabled(selectedPoints.isEmpty && boundingBoxes.isEmpty)\n Button(\"Reset\", action: resetAll)\n .padding(.trailing, 5)\n .disabled(selectedPoints.isEmpty && boundingBoxes.isEmpty)\n \n \n }\n }\n .transition(.move(edge: .top))\n }\n }\n \n private func newMask() {\n \n }\n\n private func resetAll() {\n selectedPoints.removeAll()\n boundingBoxes.removeAll()\n segmentationImages = []\n currentSegmentation = nil\n }\n \n private func undo() {\n if let lastPoint = selectedPoints.last, let lastBox = boundingBoxes.last {\n if lastPoint.dateAdded > lastBox.dateAdded {\n selectedPoints.removeLast()\n } else {\n boundingBoxes.removeLast()\n }\n } else if !selectedPoints.isEmpty {\n selectedPoints.removeLast()\n } else if !boundingBoxes.isEmpty {\n boundingBoxes.removeLast()\n }\n\n if selectedPoints.isEmpty && boundingBoxes.isEmpty {\n currentSegmentation = nil\n }\n }\n}\n\n#Preview {\n ContentView()\n}\n\n" }, { "path": "SAM2-Demo/Views/ZoomableScrollView.swift", "content": "//\n// ZoomableScrollView.swift\n// SAM2-Demo\n//\n// Created by Cyril Zakka on 9/12/24.\n//\n\nimport AppKit\nimport SwiftUI\n\nstruct ZoomableScrollView: NSViewRepresentable {\n @Binding var visibleRect: CGRect\n private var content: Content\n\n init(visibleRect: Binding, @ViewBuilder content: () -> Content) {\n self._visibleRect = visibleRect\n self.content = content()\n }\n\n func makeNSView(context: Context) -> NSScrollView {\n let scrollView = NSScrollView()\n scrollView.hasVerticalScroller = true\n scrollView.hasHorizontalScroller = true\n scrollView.autohidesScrollers = true\n scrollView.allowsMagnification = false\n scrollView.maxMagnification = 20\n scrollView.minMagnification = 1\n\n let hostedView = context.coordinator.hostingView\n hostedView.translatesAutoresizingMaskIntoConstraints = true\n hostedView.autoresizingMask = [.width, .height]\n hostedView.frame = scrollView.bounds\n scrollView.documentView = hostedView\n\n return scrollView\n }\n\n func makeCoordinator() -> Coordinator {\n let coordinator = Coordinator(hostingView: NSHostingView(rootView: self.content), parent: self)\n coordinator.listen()\n return coordinator\n }\n\n func updateNSView(_ nsView: NSScrollView, context: Context) {\n context.coordinator.hostingView.rootView = self.content\n }\n\n // MARK: - Coordinator\n\n class Coordinator: NSObject {\n var hostingView: NSHostingView\n var parent: ZoomableScrollView\n\n init(hostingView: NSHostingView, parent: ZoomableScrollView) {\n self.hostingView = hostingView\n self.parent = parent\n }\n\n func listen() {\n NotificationCenter.default.addObserver(forName: NSScrollView.didEndLiveMagnifyNotification, object: nil, queue: nil) { notification in\n let scrollView = notification.object as! NSScrollView\n print(\"did magnify: \\(scrollView.magnification), \\(scrollView.documentVisibleRect)\")\n self.parent.visibleRect = scrollView.documentVisibleRect\n }\n NotificationCenter.default.addObserver(forName: NSScrollView.didEndLiveScrollNotification, object: nil, queue: nil) { notification in\n let scrollView = notification.object as! NSScrollView\n print(\"did scroll: \\(scrollView.magnification), \\(scrollView.documentVisibleRect)\")\n self.parent.visibleRect = scrollView.documentVisibleRect\n }\n }\n }\n}\n" }, { "path": "SAM2-Demo.xcodeproj/project.pbxproj", "content": "// !$*UTF8*$!\n{\n\tarchiveVersion = 1;\n\tclasses = {\n\t};\n\tobjectVersion = 70;\n\tobjects = {\n\n/* Begin PBXBuildFile section */\n\t\tEBAB91282C88A05500F57B83 /* ArgumentParser in Frameworks */ = {isa = PBXBuildFile; productRef = EBAB91272C88A05500F57B83 /* ArgumentParser */; };\n/* End PBXBuildFile section */\n\n/* Begin PBXCopyFilesBuildPhase section */\n\t\tEBAB911D2C889D5200F57B83 /* CopyFiles */ = {\n\t\t\tisa = PBXCopyFilesBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tdstPath = /usr/share/man/man1/;\n\t\t\tdstSubfolderSpec = 0;\n\t\t\tfiles = (\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 1;\n\t\t};\n/* End PBXCopyFilesBuildPhase section */\n\n/* Begin PBXFileReference section */\n\t\tEBAB911F2C889D5200F57B83 /* sam2-cli */ = {isa = PBXFileReference; explicitFileType = \"compiled.mach-o.executable\"; includeInIndex = 0; path = \"sam2-cli\"; sourceTree = BUILT_PRODUCTS_DIR; };\n\t\tF136320F2C73AE78009DEF15 /* SAM 2 Studio.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = \"SAM 2 Studio.app\"; sourceTree = BUILT_PRODUCTS_DIR; };\n/* End PBXFileReference section */\n\n/* Begin PBXFileSystemSynchronizedBuildFileExceptionSet section */\n\t\tEBAB912A2C88A21E00F57B83 /* Exceptions for \"SAM2-Demo\" folder in \"sam2-cli\" target */ = {\n\t\t\tisa = PBXFileSystemSynchronizedBuildFileExceptionSet;\n\t\t\tmembershipExceptions = (\n\t\t\t\t\"Common/CGImage+Extension.swift\",\n\t\t\t\t\"Common/CGImage+RawBytes.swift\",\n\t\t\t\tCommon/CoreImageExtensions.swift,\n\t\t\t\tCommon/DirectoryDocument.swift,\n\t\t\t\t\"Common/MLMultiArray+Image.swift\",\n\t\t\t\tCommon/Models.swift,\n\t\t\t\tCommon/SAM2.swift,\n\t\t\t\tSAM2_1SmallImageEncoderFLOAT16.mlpackage,\n\t\t\t\tSAM2_1SmallMaskDecoderFLOAT16.mlpackage,\n\t\t\t\tSAM2_1SmallPromptEncoderFLOAT16.mlpackage,\n\t\t\t);\n\t\t\ttarget = EBAB911E2C889D5200F57B83 /* sam2-cli */;\n\t\t};\n/* End PBXFileSystemSynchronizedBuildFileExceptionSet section */\n\n/* Begin PBXFileSystemSynchronizedRootGroup section */\n\t\tEBAB91202C889D5200F57B83 /* sam2-cli */ = {\n\t\t\tisa = PBXFileSystemSynchronizedRootGroup;\n\t\t\tpath = \"sam2-cli\";\n\t\t\tsourceTree = \"\";\n\t\t};\n\t\tF13632112C73AE78009DEF15 /* SAM2-Demo */ = {\n\t\t\tisa = PBXFileSystemSynchronizedRootGroup;\n\t\t\texceptions = (\n\t\t\t\tEBAB912A2C88A21E00F57B83 /* Exceptions for \"SAM2-Demo\" folder in \"sam2-cli\" target */,\n\t\t\t);\n\t\t\tpath = \"SAM2-Demo\";\n\t\t\tsourceTree = \"\";\n\t\t};\n/* End PBXFileSystemSynchronizedRootGroup section */\n\n/* Begin PBXFrameworksBuildPhase section */\n\t\tEBAB911C2C889D5200F57B83 /* Frameworks */ = {\n\t\t\tisa = PBXFrameworksBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tfiles = (\n\t\t\t\tEBAB91282C88A05500F57B83 /* ArgumentParser in Frameworks */,\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 0;\n\t\t};\n\t\tF136320C2C73AE78009DEF15 /* Frameworks */ = {\n\t\t\tisa = PBXFrameworksBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tfiles = (\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 0;\n\t\t};\n/* End PBXFrameworksBuildPhase section */\n\n/* Begin PBXGroup section */\n\t\tF13632062C73AE77009DEF15 = {\n\t\t\tisa = PBXGroup;\n\t\t\tchildren = (\n\t\t\t\tF13632112C73AE78009DEF15 /* SAM2-Demo */,\n\t\t\t\tEBAB91202C889D5200F57B83 /* sam2-cli */,\n\t\t\t\tF13632102C73AE78009DEF15 /* Products */,\n\t\t\t);\n\t\t\tsourceTree = \"\";\n\t\t};\n\t\tF13632102C73AE78009DEF15 /* Products */ = {\n\t\t\tisa = PBXGroup;\n\t\t\tchildren = (\n\t\t\t\tF136320F2C73AE78009DEF15 /* SAM 2 Studio.app */,\n\t\t\t\tEBAB911F2C889D5200F57B83 /* sam2-cli */,\n\t\t\t);\n\t\t\tname = Products;\n\t\t\tsourceTree = \"\";\n\t\t};\n/* End PBXGroup section */\n\n/* Begin PBXNativeTarget section */\n\t\tEBAB911E2C889D5200F57B83 /* sam2-cli */ = {\n\t\t\tisa = PBXNativeTarget;\n\t\t\tbuildConfigurationList = EBAB91252C889D5200F57B83 /* Build configuration list for PBXNativeTarget \"sam2-cli\" */;\n\t\t\tbuildPhases = (\n\t\t\t\tEBAB911B2C889D5200F57B83 /* Sources */,\n\t\t\t\tEBAB911C2C889D5200F57B83 /* Frameworks */,\n\t\t\t\tEBAB911D2C889D5200F57B83 /* CopyFiles */,\n\t\t\t);\n\t\t\tbuildRules = (\n\t\t\t);\n\t\t\tdependencies = (\n\t\t\t);\n\t\t\tfileSystemSynchronizedGroups = (\n\t\t\t\tEBAB91202C889D5200F57B83 /* sam2-cli */,\n\t\t\t);\n\t\t\tname = \"sam2-cli\";\n\t\t\tpackageProductDependencies = (\n\t\t\t\tEBAB91272C88A05500F57B83 /* ArgumentParser */,\n\t\t\t);\n\t\t\tproductName = \"sam2-cli\";\n\t\t\tproductReference = EBAB911F2C889D5200F57B83 /* sam2-cli */;\n\t\t\tproductType = \"com.apple.product-type.tool\";\n\t\t};\n\t\tF136320E2C73AE78009DEF15 /* SAM2-Demo */ = {\n\t\t\tisa = PBXNativeTarget;\n\t\t\tbuildConfigurationList = F136321E2C73AE79009DEF15 /* Build configuration list for PBXNativeTarget \"SAM2-Demo\" */;\n\t\t\tbuildPhases = (\n\t\t\t\tF136320B2C73AE78009DEF15 /* Sources */,\n\t\t\t\tF136320C2C73AE78009DEF15 /* Frameworks */,\n\t\t\t\tF136320D2C73AE78009DEF15 /* Resources */,\n\t\t\t);\n\t\t\tbuildRules = (\n\t\t\t);\n\t\t\tdependencies = (\n\t\t\t);\n\t\t\tfileSystemSynchronizedGroups = (\n\t\t\t\tF13632112C73AE78009DEF15 /* SAM2-Demo */,\n\t\t\t);\n\t\t\tname = \"SAM2-Demo\";\n\t\t\tpackageProductDependencies = (\n\t\t\t);\n\t\t\tproductName = \"SAM2-Demo\";\n\t\t\tproductReference = F136320F2C73AE78009DEF15 /* SAM 2 Studio.app */;\n\t\t\tproductType = \"com.apple.product-type.application\";\n\t\t};\n/* End PBXNativeTarget section */\n\n/* Begin PBXProject section */\n\t\tF13632072C73AE77009DEF15 /* Project object */ = {\n\t\t\tisa = PBXProject;\n\t\t\tattributes = {\n\t\t\t\tBuildIndependentTargetsInParallel = 1;\n\t\t\t\tLastSwiftUpdateCheck = 1600;\n\t\t\t\tLastUpgradeCheck = 1600;\n\t\t\t\tTargetAttributes = {\n\t\t\t\t\tEBAB911E2C889D5200F57B83 = {\n\t\t\t\t\t\tCreatedOnToolsVersion = 16.0;\n\t\t\t\t\t};\n\t\t\t\t\tF136320E2C73AE78009DEF15 = {\n\t\t\t\t\t\tCreatedOnToolsVersion = 16.0;\n\t\t\t\t\t};\n\t\t\t\t};\n\t\t\t};\n\t\t\tbuildConfigurationList = F136320A2C73AE77009DEF15 /* Build configuration list for PBXProject \"SAM2-Demo\" */;\n\t\t\tdevelopmentRegion = en;\n\t\t\thasScannedForEncodings = 0;\n\t\t\tknownRegions = (\n\t\t\t\ten,\n\t\t\t\tBase,\n\t\t\t);\n\t\t\tmainGroup = F13632062C73AE77009DEF15;\n\t\t\tminimizedProjectReferenceProxies = 1;\n\t\t\tpackageReferences = (\n\t\t\t\tEBAB91262C88A05500F57B83 /* XCRemoteSwiftPackageReference \"swift-argument-parser\" */,\n\t\t\t);\n\t\t\tpreferredProjectObjectVersion = 77;\n\t\t\tproductRefGroup = F13632102C73AE78009DEF15 /* Products */;\n\t\t\tprojectDirPath = \"\";\n\t\t\tprojectRoot = \"\";\n\t\t\ttargets = (\n\t\t\t\tF136320E2C73AE78009DEF15 /* SAM2-Demo */,\n\t\t\t\tEBAB911E2C889D5200F57B83 /* sam2-cli */,\n\t\t\t);\n\t\t};\n/* End PBXProject section */\n\n/* Begin PBXResourcesBuildPhase section */\n\t\tF136320D2C73AE78009DEF15 /* Resources */ = {\n\t\t\tisa = PBXResourcesBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tfiles = (\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 0;\n\t\t};\n/* End PBXResourcesBuildPhase section */\n\n/* Begin PBXSourcesBuildPhase section */\n\t\tEBAB911B2C889D5200F57B83 /* Sources */ = {\n\t\t\tisa = PBXSourcesBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tfiles = (\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 0;\n\t\t};\n\t\tF136320B2C73AE78009DEF15 /* Sources */ = {\n\t\t\tisa = PBXSourcesBuildPhase;\n\t\t\tbuildActionMask = 2147483647;\n\t\t\tfiles = (\n\t\t\t);\n\t\t\trunOnlyForDeploymentPostprocessing = 0;\n\t\t};\n/* End PBXSourcesBuildPhase section */\n\n/* Begin XCBuildConfiguration section */\n\t\tEBAB91232C889D5200F57B83 /* Debug */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tCODE_SIGN_STYLE = Automatic;\n\t\t\t\tDEVELOPMENT_TEAM = \"\";\n\t\t\t\tENABLE_HARDENED_RUNTIME = YES;\n\t\t\t\tPRODUCT_NAME = \"$(TARGET_NAME)\";\n\t\t\t\tSWIFT_VERSION = 5.0;\n\t\t\t};\n\t\t\tname = Debug;\n\t\t};\n\t\tEBAB91242C889D5200F57B83 /* Release */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tCODE_SIGN_STYLE = Automatic;\n\t\t\t\tDEVELOPMENT_TEAM = \"\";\n\t\t\t\tENABLE_HARDENED_RUNTIME = YES;\n\t\t\t\tPRODUCT_NAME = \"$(TARGET_NAME)\";\n\t\t\t\tSWIFT_VERSION = 5.0;\n\t\t\t};\n\t\t\tname = Release;\n\t\t};\n\t\tF136321C2C73AE79009DEF15 /* Debug */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tALWAYS_SEARCH_USER_PATHS = NO;\n\t\t\t\tARCHS = arm64;\n\t\t\t\tASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;\n\t\t\t\tCLANG_ANALYZER_NONNULL = YES;\n\t\t\t\tCLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;\n\t\t\t\tCLANG_CXX_LANGUAGE_STANDARD = \"gnu++20\";\n\t\t\t\tCLANG_ENABLE_MODULES = YES;\n\t\t\t\tCLANG_ENABLE_OBJC_ARC = YES;\n\t\t\t\tCLANG_ENABLE_OBJC_WEAK = YES;\n\t\t\t\tCLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;\n\t\t\t\tCLANG_WARN_BOOL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_COMMA = YES;\n\t\t\t\tCLANG_WARN_CONSTANT_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;\n\t\t\t\tCLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;\n\t\t\t\tCLANG_WARN_DOCUMENTATION_COMMENTS = YES;\n\t\t\t\tCLANG_WARN_EMPTY_BODY = YES;\n\t\t\t\tCLANG_WARN_ENUM_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_INFINITE_RECURSION = YES;\n\t\t\t\tCLANG_WARN_INT_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;\n\t\t\t\tCLANG_WARN_OBJC_LITERAL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;\n\t\t\t\tCLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;\n\t\t\t\tCLANG_WARN_RANGE_LOOP_ANALYSIS = YES;\n\t\t\t\tCLANG_WARN_STRICT_PROTOTYPES = YES;\n\t\t\t\tCLANG_WARN_SUSPICIOUS_MOVE = YES;\n\t\t\t\tCLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;\n\t\t\t\tCLANG_WARN_UNREACHABLE_CODE = YES;\n\t\t\t\tCLANG_WARN__DUPLICATE_METHOD_MATCH = YES;\n\t\t\t\tCOPY_PHASE_STRIP = NO;\n\t\t\t\tDEBUG_INFORMATION_FORMAT = dwarf;\n\t\t\t\tDEVELOPMENT_TEAM = 2EADP68M95;\n\t\t\t\tENABLE_STRICT_OBJC_MSGSEND = YES;\n\t\t\t\tENABLE_TESTABILITY = YES;\n\t\t\t\tENABLE_USER_SCRIPT_SANDBOXING = YES;\n\t\t\t\tGCC_C_LANGUAGE_STANDARD = gnu17;\n\t\t\t\tGCC_DYNAMIC_NO_PIC = NO;\n\t\t\t\tGCC_NO_COMMON_BLOCKS = YES;\n\t\t\t\tGCC_OPTIMIZATION_LEVEL = 0;\n\t\t\t\tGCC_PREPROCESSOR_DEFINITIONS = (\n\t\t\t\t\t\"DEBUG=1\",\n\t\t\t\t\t\"$(inherited)\",\n\t\t\t\t);\n\t\t\t\tGCC_WARN_64_TO_32_BIT_CONVERSION = YES;\n\t\t\t\tGCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;\n\t\t\t\tGCC_WARN_UNDECLARED_SELECTOR = YES;\n\t\t\t\tGCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;\n\t\t\t\tGCC_WARN_UNUSED_FUNCTION = YES;\n\t\t\t\tGCC_WARN_UNUSED_VARIABLE = YES;\n\t\t\t\tLOCALIZATION_PREFERS_STRING_CATALOGS = YES;\n\t\t\t\tMACOSX_DEPLOYMENT_TARGET = 14.3;\n\t\t\t\tMTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;\n\t\t\t\tMTL_FAST_MATH = YES;\n\t\t\t\tONLY_ACTIVE_ARCH = YES;\n\t\t\t\tSDKROOT = macosx;\n\t\t\t\tSWIFT_ACTIVE_COMPILATION_CONDITIONS = \"DEBUG $(inherited)\";\n\t\t\t\tSWIFT_OPTIMIZATION_LEVEL = \"-Onone\";\n\t\t\t};\n\t\t\tname = Debug;\n\t\t};\n\t\tF136321D2C73AE79009DEF15 /* Release */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tALWAYS_SEARCH_USER_PATHS = NO;\n\t\t\t\tARCHS = arm64;\n\t\t\t\tASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;\n\t\t\t\tCLANG_ANALYZER_NONNULL = YES;\n\t\t\t\tCLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;\n\t\t\t\tCLANG_CXX_LANGUAGE_STANDARD = \"gnu++20\";\n\t\t\t\tCLANG_ENABLE_MODULES = YES;\n\t\t\t\tCLANG_ENABLE_OBJC_ARC = YES;\n\t\t\t\tCLANG_ENABLE_OBJC_WEAK = YES;\n\t\t\t\tCLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;\n\t\t\t\tCLANG_WARN_BOOL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_COMMA = YES;\n\t\t\t\tCLANG_WARN_CONSTANT_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;\n\t\t\t\tCLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;\n\t\t\t\tCLANG_WARN_DOCUMENTATION_COMMENTS = YES;\n\t\t\t\tCLANG_WARN_EMPTY_BODY = YES;\n\t\t\t\tCLANG_WARN_ENUM_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_INFINITE_RECURSION = YES;\n\t\t\t\tCLANG_WARN_INT_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;\n\t\t\t\tCLANG_WARN_OBJC_LITERAL_CONVERSION = YES;\n\t\t\t\tCLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;\n\t\t\t\tCLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;\n\t\t\t\tCLANG_WARN_RANGE_LOOP_ANALYSIS = YES;\n\t\t\t\tCLANG_WARN_STRICT_PROTOTYPES = YES;\n\t\t\t\tCLANG_WARN_SUSPICIOUS_MOVE = YES;\n\t\t\t\tCLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;\n\t\t\t\tCLANG_WARN_UNREACHABLE_CODE = YES;\n\t\t\t\tCLANG_WARN__DUPLICATE_METHOD_MATCH = YES;\n\t\t\t\tCOPY_PHASE_STRIP = NO;\n\t\t\t\tDEBUG_INFORMATION_FORMAT = \"dwarf-with-dsym\";\n\t\t\t\tDEVELOPMENT_TEAM = 2EADP68M95;\n\t\t\t\tENABLE_NS_ASSERTIONS = NO;\n\t\t\t\tENABLE_STRICT_OBJC_MSGSEND = YES;\n\t\t\t\tENABLE_USER_SCRIPT_SANDBOXING = YES;\n\t\t\t\tGCC_C_LANGUAGE_STANDARD = gnu17;\n\t\t\t\tGCC_NO_COMMON_BLOCKS = YES;\n\t\t\t\tGCC_WARN_64_TO_32_BIT_CONVERSION = YES;\n\t\t\t\tGCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;\n\t\t\t\tGCC_WARN_UNDECLARED_SELECTOR = YES;\n\t\t\t\tGCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;\n\t\t\t\tGCC_WARN_UNUSED_FUNCTION = YES;\n\t\t\t\tGCC_WARN_UNUSED_VARIABLE = YES;\n\t\t\t\tLOCALIZATION_PREFERS_STRING_CATALOGS = YES;\n\t\t\t\tMACOSX_DEPLOYMENT_TARGET = 14.3;\n\t\t\t\tMTL_ENABLE_DEBUG_INFO = NO;\n\t\t\t\tMTL_FAST_MATH = YES;\n\t\t\t\tSDKROOT = macosx;\n\t\t\t\tSWIFT_COMPILATION_MODE = wholemodule;\n\t\t\t};\n\t\t\tname = Release;\n\t\t};\n\t\tF136321F2C73AE79009DEF15 /* Debug */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;\n\t\t\t\tASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;\n\t\t\t\tCODE_SIGN_ENTITLEMENTS = \"SAM2-Demo/SAM2_Demo.entitlements\";\n\t\t\t\tCODE_SIGN_STYLE = Automatic;\n\t\t\t\tCOMBINE_HIDPI_IMAGES = YES;\n\t\t\t\tCURRENT_PROJECT_VERSION = 1;\n\t\t\t\tDEVELOPMENT_ASSET_PATHS = \"\\\"SAM2-Demo/Preview Content\\\"\";\n\t\t\t\tDEVELOPMENT_TEAM = \"\";\n\t\t\t\tENABLE_HARDENED_RUNTIME = YES;\n\t\t\t\tENABLE_PREVIEWS = YES;\n\t\t\t\tGENERATE_INFOPLIST_FILE = YES;\n\t\t\t\tINFOPLIST_KEY_NSHumanReadableCopyright = \"\";\n\t\t\t\tLD_RUNPATH_SEARCH_PATHS = (\n\t\t\t\t\t\"$(inherited)\",\n\t\t\t\t\t\"@executable_path/../Frameworks\",\n\t\t\t\t);\n\t\t\t\tMARKETING_VERSION = 1.0;\n\t\t\t\tPRODUCT_BUNDLE_IDENTIFIER = \"co.huggingface.sam-2-studio\";\n\t\t\t\tPRODUCT_NAME = \"SAM 2 Studio\";\n\t\t\t\tSWIFT_EMIT_LOC_STRINGS = YES;\n\t\t\t\tSWIFT_VERSION = 5.0;\n\t\t\t};\n\t\t\tname = Debug;\n\t\t};\n\t\tF13632202C73AE79009DEF15 /* Release */ = {\n\t\t\tisa = XCBuildConfiguration;\n\t\t\tbuildSettings = {\n\t\t\t\tASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;\n\t\t\t\tASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;\n\t\t\t\tCODE_SIGN_ENTITLEMENTS = \"SAM2-Demo/SAM2_Demo.entitlements\";\n\t\t\t\tCODE_SIGN_STYLE = Automatic;\n\t\t\t\tCOMBINE_HIDPI_IMAGES = YES;\n\t\t\t\tCURRENT_PROJECT_VERSION = 1;\n\t\t\t\tDEVELOPMENT_ASSET_PATHS = \"\\\"SAM2-Demo/Preview Content\\\"\";\n\t\t\t\tDEVELOPMENT_TEAM = \"\";\n\t\t\t\tENABLE_HARDENED_RUNTIME = YES;\n\t\t\t\tENABLE_PREVIEWS = YES;\n\t\t\t\tGENERATE_INFOPLIST_FILE = YES;\n\t\t\t\tINFOPLIST_KEY_NSHumanReadableCopyright = \"\";\n\t\t\t\tLD_RUNPATH_SEARCH_PATHS = (\n\t\t\t\t\t\"$(inherited)\",\n\t\t\t\t\t\"@executable_path/../Frameworks\",\n\t\t\t\t);\n\t\t\t\tMARKETING_VERSION = 1.0;\n\t\t\t\tPRODUCT_BUNDLE_IDENTIFIER = \"co.huggingface.sam-2-studio\";\n\t\t\t\tPRODUCT_NAME = \"SAM 2 Studio\";\n\t\t\t\tSWIFT_EMIT_LOC_STRINGS = YES;\n\t\t\t\tSWIFT_VERSION = 5.0;\n\t\t\t};\n\t\t\tname = Release;\n\t\t};\n/* End XCBuildConfiguration section */\n\n/* Begin XCConfigurationList section */\n\t\tEBAB91252C889D5200F57B83 /* Build configuration list for PBXNativeTarget \"sam2-cli\" */ = {\n\t\t\tisa = XCConfigurationList;\n\t\t\tbuildConfigurations = (\n\t\t\t\tEBAB91232C889D5200F57B83 /* Debug */,\n\t\t\t\tEBAB91242C889D5200F57B83 /* Release */,\n\t\t\t);\n\t\t\tdefaultConfigurationIsVisible = 0;\n\t\t\tdefaultConfigurationName = Release;\n\t\t};\n\t\tF136320A2C73AE77009DEF15 /* Build configuration list for PBXProject \"SAM2-Demo\" */ = {\n\t\t\tisa = XCConfigurationList;\n\t\t\tbuildConfigurations = (\n\t\t\t\tF136321C2C73AE79009DEF15 /* Debug */,\n\t\t\t\tF136321D2C73AE79009DEF15 /* Release */,\n\t\t\t);\n\t\t\tdefaultConfigurationIsVisible = 0;\n\t\t\tdefaultConfigurationName = Release;\n\t\t};\n\t\tF136321E2C73AE79009DEF15 /* Build configuration list for PBXNativeTarget \"SAM2-Demo\" */ = {\n\t\t\tisa = XCConfigurationList;\n\t\t\tbuildConfigurations = (\n\t\t\t\tF136321F2C73AE79009DEF15 /* Debug */,\n\t\t\t\tF13632202C73AE79009DEF15 /* Release */,\n\t\t\t);\n\t\t\tdefaultConfigurationIsVisible = 0;\n\t\t\tdefaultConfigurationName = Release;\n\t\t};\n/* End XCConfigurationList section */\n\n/* Begin XCRemoteSwiftPackageReference section */\n\t\tEBAB91262C88A05500F57B83 /* XCRemoteSwiftPackageReference \"swift-argument-parser\" */ = {\n\t\t\tisa = XCRemoteSwiftPackageReference;\n\t\t\trepositoryURL = \"https://github.com/apple/swift-argument-parser.git\";\n\t\t\trequirement = {\n\t\t\t\tkind = upToNextMajorVersion;\n\t\t\t\tminimumVersion = 1.5.0;\n\t\t\t};\n\t\t};\n/* End XCRemoteSwiftPackageReference section */\n\n/* Begin XCSwiftPackageProductDependency section */\n\t\tEBAB91272C88A05500F57B83 /* ArgumentParser */ = {\n\t\t\tisa = XCSwiftPackageProductDependency;\n\t\t\tpackage = EBAB91262C88A05500F57B83 /* XCRemoteSwiftPackageReference \"swift-argument-parser\" */;\n\t\t\tproductName = ArgumentParser;\n\t\t};\n/* End XCSwiftPackageProductDependency section */\n\t};\n\trootObject = F13632072C73AE77009DEF15 /* Project object */;\n}\n" }, { "path": "SAM2-Demo.xcodeproj/project.xcworkspace/contents.xcworkspacedata", "content": "\n\n \n \n\n" }, { "path": "SAM2-Demo.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved", "content": "{\n \"originHash\" : \"59ba1edda695b389d6c9ac1809891cd779e4024f505b0ce1a9d5202b6762e38a\",\n \"pins\" : [\n {\n \"identity\" : \"swift-argument-parser\",\n \"kind\" : \"remoteSourceControl\",\n \"location\" : \"https://github.com/apple/swift-argument-parser.git\",\n \"state\" : {\n \"revision\" : \"41982a3656a71c768319979febd796c6fd111d5c\",\n \"version\" : \"1.5.0\"\n }\n }\n ],\n \"version\" : 3\n}\n" }, { "path": "sam2-cli/MainCommand.swift", "content": "import ArgumentParser\nimport CoreImage\nimport CoreML\nimport ImageIO\nimport UniformTypeIdentifiers\nimport Combine\n\nlet context = CIContext(options: [.outputColorSpace: NSNull()])\n\nenum PointType: Int, ExpressibleByArgument {\n case background = 0\n case foreground = 1\n\n var asCategory: SAMCategory {\n switch self {\n case .background:\n return SAMCategory.background\n case .foreground:\n return SAMCategory.foreground\n }\n }\n}\n\n@main\nstruct MainCommand: AsyncParsableCommand {\n static let configuration = CommandConfiguration(\n commandName: \"sam2-cli\",\n abstract: \"Perform segmentation using the SAM v2 model.\"\n )\n\n @Option(name: .shortAndLong, help: \"The input image file.\")\n var input: String\n\n // TODO: multiple points\n @Option(name: .shortAndLong, parsing: .upToNextOption, help: \"List of input coordinates in format 'x,y'. Coordinates are relative to the input image size. Separate multiple entries with spaces, but don't use spaces between the coordinates.\")\n var points: [CGPoint]\n\n @Option(name: .shortAndLong, parsing: .upToNextOption, help: \"Point types that correspond to the input points. Use as many as points, 0 for background and 1 for foreground.\")\n var types: [PointType]\n\n @Option(name: .shortAndLong, help: \"The output PNG image file, showing the segmentation map overlaid on top of the original image.\")\n var output: String\n\n @Option(name: [.long, .customShort(\"k\")], help: \"The output file name for the segmentation mask.\")\n var mask: String? = nil\n\n @MainActor\n mutating func run() async throws {\n // TODO: specify directory with loadable .mlpackages instead\n let sam = try await SAM2.load()\n print(\"Models loaded in: \\(String(describing: sam.initializationTime))\")\n let targetSize = sam.inputSize\n\n // Load the input image\n guard let inputImage = CIImage(contentsOf: URL(filePath: input), options: [.colorSpace: NSNull()]) else {\n print(\"Failed to load image.\")\n throw ExitCode(EXIT_FAILURE)\n }\n print(\"Original image size \\(inputImage.extent)\")\n\n // Resize the image to match the model's expected input\n let resizedImage = inputImage.resized(to: targetSize)\n\n // Convert to a pixel buffer\n guard let pixelBuffer = context.render(resizedImage, pixelFormat: kCVPixelFormatType_32ARGB) else {\n print(\"Failed to create pixel buffer for input image.\")\n throw ExitCode(EXIT_FAILURE)\n }\n\n // Execute the model\n let clock = ContinuousClock()\n let start = clock.now\n try await sam.getImageEncoding(from: pixelBuffer)\n let duration = clock.now - start\n print(\"Image encoding took \\(duration.formatted(.units(allowed: [.seconds, .milliseconds])))\")\n\n let startMask = clock.now\n let pointSequence = zip(points, types).map { point, type in\n SAMPoint(coordinates:point, category:type.asCategory)\n }\n try await sam.getPromptEncoding(from: pointSequence, with: inputImage.extent.size)\n guard let maskImage = try await sam.getMask(for: inputImage.extent.size) else {\n throw ExitCode(EXIT_FAILURE)\n }\n let maskDuration = clock.now - startMask\n print(\"Prompt encoding and mask generation took \\(maskDuration.formatted(.units(allowed: [.seconds, .milliseconds])))\")\n\n // Write masks\n if let mask = mask {\n context.writePNG(maskImage, to: URL(filePath: mask))\n }\n\n // Overlay over original and save\n guard let outputImage = maskImage.withAlpha(0.6)?.composited(over: inputImage) else {\n print(\"Failed to blend mask.\")\n throw ExitCode(EXIT_FAILURE)\n }\n context.writePNG(outputImage, to: URL(filePath: output))\n }\n}\n\nextension CGPoint: ExpressibleByArgument {\n public init?(argument: String) {\n let components = argument.split(separator: \",\").map(String.init)\n\n guard components.count == 2,\n let x = Double(components[0]),\n let y = Double(components[1]) else {\n return nil\n }\n\n self.init(x: x, y: y)\n }\n}\n" } ]