├── .gitignore ├── BackgroundRemovalWithCoreMLSample.xcodeproj ├── project.pbxproj └── project.xcworkspace │ ├── contents.xcworkspacedata │ └── xcshareddata │ └── IDEWorkspaceChecks.plist ├── BackgroundRemovalWithCoreMLSample ├── AppDelegate.swift ├── Assets.xcassets │ ├── AppIcon.appiconset │ │ └── Contents.json │ ├── Contents.json │ └── cat.imageset │ │ ├── 5122257268_42d3a32729_z.jpg │ │ └── Contents.json ├── Base.lproj │ ├── LaunchScreen.storyboard │ └── Main.storyboard ├── CoreMLHelpers │ ├── Array+Extensions.swift │ ├── CGImage+CVPixelBuffer.swift │ ├── CGImage+RawBytes.swift │ ├── CGImagePropertyOrientation.swift │ ├── CVPixelBuffer+Helpers.swift │ ├── LICENSE.txt │ ├── MLMultiArray+Helpers.swift │ ├── MLMultiArray+Image.swift │ ├── Math.swift │ ├── NonMaxSuppression.swift │ ├── Predictions.swift │ ├── UIImage+CVPixelBuffer.swift │ ├── UIImage+Extensions.swift │ └── UIImage+RawBytes.swift ├── DeepLabV3.mlmodel ├── Info.plist └── ViewController.swift ├── LICENSE ├── README.md └── ScreenShot.png /.gitignore: -------------------------------------------------------------------------------- 1 | # Xcode 2 | # 3 | # gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore 4 | 5 | ## Build generated 6 | build/ 7 | DerivedData/ 8 | 9 | ## Various settings 10 | *.pbxuser 11 | !default.pbxuser 12 | *.mode1v3 13 | !default.mode1v3 14 | *.mode2v3 15 | !default.mode2v3 16 | *.perspectivev3 17 | !default.perspectivev3 18 | xcuserdata/ 19 | 20 | ## Other 21 | *.moved-aside 22 | *.xccheckout 23 | *.xcscmblueprint 24 | 25 | ## Obj-C/Swift specific 26 | *.hmap 27 | *.ipa 28 | *.dSYM.zip 29 | *.dSYM 30 | 31 | ## Playgrounds 32 | timeline.xctimeline 33 | playground.xcworkspace 34 | 35 | # Swift Package Manager 36 | # 37 | # Add this line if you want to avoid checking in source code from Swift Package Manager dependencies. 38 | # Packages/ 39 | # Package.pins 40 | # Package.resolved 41 | .build/ 42 | 43 | # CocoaPods 44 | # 45 | # We recommend against adding the Pods directory to your .gitignore. However 46 | # you should judge for yourself, the pros and cons are mentioned at: 47 | # https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control 48 | # 49 | # Pods/ 50 | 51 | # Carthage 52 | # 53 | # Add this line if you want to avoid checking in source code from Carthage dependencies. 54 | # Carthage/Checkouts 55 | 56 | Carthage/Build 57 | 58 | # fastlane 59 | # 60 | # It is recommended to not store the screenshots in the git repo. Instead, use fastlane to re-generate the 61 | # screenshots whenever they are needed. 62 | # For more information about the recommended setup visit: 63 | # https://docs.fastlane.tools/best-practices/source-control/#source-control 64 | 65 | fastlane/report.xml 66 | fastlane/Preview.html 67 | fastlane/screenshots/**/*.png 68 | fastlane/test_output 69 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample.xcodeproj/project.pbxproj: -------------------------------------------------------------------------------- 1 | // !$*UTF8*$! 2 | { 3 | archiveVersion = 1; 4 | classes = { 5 | }; 6 | objectVersion = 50; 7 | objects = { 8 | 9 | /* Begin PBXBuildFile section */ 10 | 2139408F22D32D2300B23738 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2139408E22D32D2300B23738 /* AppDelegate.swift */; }; 11 | 2139409122D32D2300B23738 /* ViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2139409022D32D2300B23738 /* ViewController.swift */; }; 12 | 2139409422D32D2300B23738 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 2139409222D32D2300B23738 /* Main.storyboard */; }; 13 | 2139409622D32D2400B23738 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 2139409522D32D2400B23738 /* Assets.xcassets */; }; 14 | 2139409922D32D2400B23738 /* LaunchScreen.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 2139409722D32D2400B23738 /* LaunchScreen.storyboard */; }; 15 | 213940A122D32DC000B23738 /* DeepLabV3.mlmodel in Sources */ = {isa = PBXBuildFile; fileRef = 213940A022D32DC000B23738 /* DeepLabV3.mlmodel */; }; 16 | 213940A722D32F5200B23738 /* LICENSE.txt in Resources */ = {isa = PBXBuildFile; fileRef = 213940A622D32F5200B23738 /* LICENSE.txt */; }; 17 | 213940B622D32F6100B23738 /* Array+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940A922D32F6100B23738 /* Array+Extensions.swift */; }; 18 | 213940B722D32F6100B23738 /* Math.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AA22D32F6100B23738 /* Math.swift */; }; 19 | 213940B822D32F6100B23738 /* UIImage+CVPixelBuffer.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AB22D32F6100B23738 /* UIImage+CVPixelBuffer.swift */; }; 20 | 213940B922D32F6100B23738 /* Predictions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AC22D32F6100B23738 /* Predictions.swift */; }; 21 | 213940BA22D32F6100B23738 /* UIImage+Extensions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AD22D32F6100B23738 /* UIImage+Extensions.swift */; }; 22 | 213940BB22D32F6100B23738 /* CGImagePropertyOrientation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AE22D32F6100B23738 /* CGImagePropertyOrientation.swift */; }; 23 | 213940BC22D32F6100B23738 /* UIImage+RawBytes.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940AF22D32F6100B23738 /* UIImage+RawBytes.swift */; }; 24 | 213940BD22D32F6100B23738 /* MLMultiArray+Helpers.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B022D32F6100B23738 /* MLMultiArray+Helpers.swift */; }; 25 | 213940BE22D32F6100B23738 /* CVPixelBuffer+Helpers.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B122D32F6100B23738 /* CVPixelBuffer+Helpers.swift */; }; 26 | 213940BF22D32F6100B23738 /* CGImage+CVPixelBuffer.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B222D32F6100B23738 /* CGImage+CVPixelBuffer.swift */; }; 27 | 213940C022D32F6100B23738 /* NonMaxSuppression.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B322D32F6100B23738 /* NonMaxSuppression.swift */; }; 28 | 213940C122D32F6100B23738 /* CGImage+RawBytes.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B422D32F6100B23738 /* CGImage+RawBytes.swift */; }; 29 | 213940C222D32F6100B23738 /* MLMultiArray+Image.swift in Sources */ = {isa = PBXBuildFile; fileRef = 213940B522D32F6100B23738 /* MLMultiArray+Image.swift */; }; 30 | /* End PBXBuildFile section */ 31 | 32 | /* Begin PBXFileReference section */ 33 | 2139408B22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = BackgroundRemovalWithCoreMLSample.app; sourceTree = BUILT_PRODUCTS_DIR; }; 34 | 2139408E22D32D2300B23738 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = ""; }; 35 | 2139409022D32D2300B23738 /* ViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ViewController.swift; sourceTree = ""; }; 36 | 2139409322D32D2300B23738 /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/Main.storyboard; sourceTree = ""; }; 37 | 2139409522D32D2400B23738 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; 38 | 2139409822D32D2400B23738 /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/LaunchScreen.storyboard; sourceTree = ""; }; 39 | 2139409A22D32D2400B23738 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = ""; }; 40 | 213940A022D32DC000B23738 /* DeepLabV3.mlmodel */ = {isa = PBXFileReference; lastKnownFileType = file.mlmodel; path = DeepLabV3.mlmodel; sourceTree = ""; }; 41 | 213940A622D32F5200B23738 /* LICENSE.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = LICENSE.txt; sourceTree = ""; }; 42 | 213940A922D32F6100B23738 /* Array+Extensions.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "Array+Extensions.swift"; sourceTree = ""; }; 43 | 213940AA22D32F6100B23738 /* Math.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Math.swift; sourceTree = ""; }; 44 | 213940AB22D32F6100B23738 /* UIImage+CVPixelBuffer.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+CVPixelBuffer.swift"; sourceTree = ""; }; 45 | 213940AC22D32F6100B23738 /* Predictions.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Predictions.swift; sourceTree = ""; }; 46 | 213940AD22D32F6100B23738 /* UIImage+Extensions.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+Extensions.swift"; sourceTree = ""; }; 47 | 213940AE22D32F6100B23738 /* CGImagePropertyOrientation.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = CGImagePropertyOrientation.swift; sourceTree = ""; }; 48 | 213940AF22D32F6100B23738 /* UIImage+RawBytes.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+RawBytes.swift"; sourceTree = ""; }; 49 | 213940B022D32F6100B23738 /* MLMultiArray+Helpers.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "MLMultiArray+Helpers.swift"; sourceTree = ""; }; 50 | 213940B122D32F6100B23738 /* CVPixelBuffer+Helpers.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "CVPixelBuffer+Helpers.swift"; sourceTree = ""; }; 51 | 213940B222D32F6100B23738 /* CGImage+CVPixelBuffer.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "CGImage+CVPixelBuffer.swift"; sourceTree = ""; }; 52 | 213940B322D32F6100B23738 /* NonMaxSuppression.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = NonMaxSuppression.swift; sourceTree = ""; }; 53 | 213940B422D32F6100B23738 /* CGImage+RawBytes.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "CGImage+RawBytes.swift"; sourceTree = ""; }; 54 | 213940B522D32F6100B23738 /* MLMultiArray+Image.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "MLMultiArray+Image.swift"; sourceTree = ""; }; 55 | 213940C522D35FF600B23738 /* README.md */ = {isa = PBXFileReference; lastKnownFileType = net.daringfireball.markdown; path = README.md; sourceTree = ""; }; 56 | /* End PBXFileReference section */ 57 | 58 | /* Begin PBXFrameworksBuildPhase section */ 59 | 2139408822D32D2300B23738 /* Frameworks */ = { 60 | isa = PBXFrameworksBuildPhase; 61 | buildActionMask = 2147483647; 62 | files = ( 63 | ); 64 | runOnlyForDeploymentPostprocessing = 0; 65 | }; 66 | /* End PBXFrameworksBuildPhase section */ 67 | 68 | /* Begin PBXGroup section */ 69 | 2139408222D32D2300B23738 = { 70 | isa = PBXGroup; 71 | children = ( 72 | 213940C522D35FF600B23738 /* README.md */, 73 | 2139408D22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample */, 74 | 2139408C22D32D2300B23738 /* Products */, 75 | ); 76 | sourceTree = ""; 77 | }; 78 | 2139408C22D32D2300B23738 /* Products */ = { 79 | isa = PBXGroup; 80 | children = ( 81 | 2139408B22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample.app */, 82 | ); 83 | name = Products; 84 | sourceTree = ""; 85 | }; 86 | 2139408D22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample */ = { 87 | isa = PBXGroup; 88 | children = ( 89 | 213940A822D32F6100B23738 /* CoreMLHelpers */, 90 | 2139408E22D32D2300B23738 /* AppDelegate.swift */, 91 | 2139409022D32D2300B23738 /* ViewController.swift */, 92 | 213940A022D32DC000B23738 /* DeepLabV3.mlmodel */, 93 | 2139409222D32D2300B23738 /* Main.storyboard */, 94 | 2139409522D32D2400B23738 /* Assets.xcassets */, 95 | 2139409722D32D2400B23738 /* LaunchScreen.storyboard */, 96 | 2139409A22D32D2400B23738 /* Info.plist */, 97 | ); 98 | path = BackgroundRemovalWithCoreMLSample; 99 | sourceTree = ""; 100 | }; 101 | 213940A822D32F6100B23738 /* CoreMLHelpers */ = { 102 | isa = PBXGroup; 103 | children = ( 104 | 213940A622D32F5200B23738 /* LICENSE.txt */, 105 | 213940A922D32F6100B23738 /* Array+Extensions.swift */, 106 | 213940AA22D32F6100B23738 /* Math.swift */, 107 | 213940AB22D32F6100B23738 /* UIImage+CVPixelBuffer.swift */, 108 | 213940AC22D32F6100B23738 /* Predictions.swift */, 109 | 213940AD22D32F6100B23738 /* UIImage+Extensions.swift */, 110 | 213940AE22D32F6100B23738 /* CGImagePropertyOrientation.swift */, 111 | 213940AF22D32F6100B23738 /* UIImage+RawBytes.swift */, 112 | 213940B022D32F6100B23738 /* MLMultiArray+Helpers.swift */, 113 | 213940B122D32F6100B23738 /* CVPixelBuffer+Helpers.swift */, 114 | 213940B222D32F6100B23738 /* CGImage+CVPixelBuffer.swift */, 115 | 213940B322D32F6100B23738 /* NonMaxSuppression.swift */, 116 | 213940B422D32F6100B23738 /* CGImage+RawBytes.swift */, 117 | 213940B522D32F6100B23738 /* MLMultiArray+Image.swift */, 118 | ); 119 | path = CoreMLHelpers; 120 | sourceTree = ""; 121 | }; 122 | /* End PBXGroup section */ 123 | 124 | /* Begin PBXNativeTarget section */ 125 | 2139408A22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample */ = { 126 | isa = PBXNativeTarget; 127 | buildConfigurationList = 2139409D22D32D2400B23738 /* Build configuration list for PBXNativeTarget "BackgroundRemovalWithCoreMLSample" */; 128 | buildPhases = ( 129 | 2139408722D32D2300B23738 /* Sources */, 130 | 2139408822D32D2300B23738 /* Frameworks */, 131 | 2139408922D32D2300B23738 /* Resources */, 132 | ); 133 | buildRules = ( 134 | ); 135 | dependencies = ( 136 | ); 137 | name = BackgroundRemovalWithCoreMLSample; 138 | productName = BackgroundRemovalWithCoreMLSample; 139 | productReference = 2139408B22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample.app */; 140 | productType = "com.apple.product-type.application"; 141 | }; 142 | /* End PBXNativeTarget section */ 143 | 144 | /* Begin PBXProject section */ 145 | 2139408322D32D2300B23738 /* Project object */ = { 146 | isa = PBXProject; 147 | attributes = { 148 | LastSwiftUpdateCheck = 1020; 149 | LastUpgradeCheck = 1020; 150 | ORGANIZATIONNAME = "t-chen"; 151 | TargetAttributes = { 152 | 2139408A22D32D2300B23738 = { 153 | CreatedOnToolsVersion = 10.2.1; 154 | }; 155 | }; 156 | }; 157 | buildConfigurationList = 2139408622D32D2300B23738 /* Build configuration list for PBXProject "BackgroundRemovalWithCoreMLSample" */; 158 | compatibilityVersion = "Xcode 9.3"; 159 | developmentRegion = en; 160 | hasScannedForEncodings = 0; 161 | knownRegions = ( 162 | en, 163 | Base, 164 | ); 165 | mainGroup = 2139408222D32D2300B23738; 166 | productRefGroup = 2139408C22D32D2300B23738 /* Products */; 167 | projectDirPath = ""; 168 | projectRoot = ""; 169 | targets = ( 170 | 2139408A22D32D2300B23738 /* BackgroundRemovalWithCoreMLSample */, 171 | ); 172 | }; 173 | /* End PBXProject section */ 174 | 175 | /* Begin PBXResourcesBuildPhase section */ 176 | 2139408922D32D2300B23738 /* Resources */ = { 177 | isa = PBXResourcesBuildPhase; 178 | buildActionMask = 2147483647; 179 | files = ( 180 | 213940A722D32F5200B23738 /* LICENSE.txt in Resources */, 181 | 2139409922D32D2400B23738 /* LaunchScreen.storyboard in Resources */, 182 | 2139409622D32D2400B23738 /* Assets.xcassets in Resources */, 183 | 2139409422D32D2300B23738 /* Main.storyboard in Resources */, 184 | ); 185 | runOnlyForDeploymentPostprocessing = 0; 186 | }; 187 | /* End PBXResourcesBuildPhase section */ 188 | 189 | /* Begin PBXSourcesBuildPhase section */ 190 | 2139408722D32D2300B23738 /* Sources */ = { 191 | isa = PBXSourcesBuildPhase; 192 | buildActionMask = 2147483647; 193 | files = ( 194 | 213940BF22D32F6100B23738 /* CGImage+CVPixelBuffer.swift in Sources */, 195 | 213940BE22D32F6100B23738 /* CVPixelBuffer+Helpers.swift in Sources */, 196 | 213940B622D32F6100B23738 /* Array+Extensions.swift in Sources */, 197 | 213940C122D32F6100B23738 /* CGImage+RawBytes.swift in Sources */, 198 | 213940BD22D32F6100B23738 /* MLMultiArray+Helpers.swift in Sources */, 199 | 213940BB22D32F6100B23738 /* CGImagePropertyOrientation.swift in Sources */, 200 | 213940C222D32F6100B23738 /* MLMultiArray+Image.swift in Sources */, 201 | 213940B722D32F6100B23738 /* Math.swift in Sources */, 202 | 213940B922D32F6100B23738 /* Predictions.swift in Sources */, 203 | 213940BC22D32F6100B23738 /* UIImage+RawBytes.swift in Sources */, 204 | 213940A122D32DC000B23738 /* DeepLabV3.mlmodel in Sources */, 205 | 2139409122D32D2300B23738 /* ViewController.swift in Sources */, 206 | 213940BA22D32F6100B23738 /* UIImage+Extensions.swift in Sources */, 207 | 2139408F22D32D2300B23738 /* AppDelegate.swift in Sources */, 208 | 213940C022D32F6100B23738 /* NonMaxSuppression.swift in Sources */, 209 | 213940B822D32F6100B23738 /* UIImage+CVPixelBuffer.swift in Sources */, 210 | ); 211 | runOnlyForDeploymentPostprocessing = 0; 212 | }; 213 | /* End PBXSourcesBuildPhase section */ 214 | 215 | /* Begin PBXVariantGroup section */ 216 | 2139409222D32D2300B23738 /* Main.storyboard */ = { 217 | isa = PBXVariantGroup; 218 | children = ( 219 | 2139409322D32D2300B23738 /* Base */, 220 | ); 221 | name = Main.storyboard; 222 | sourceTree = ""; 223 | }; 224 | 2139409722D32D2400B23738 /* LaunchScreen.storyboard */ = { 225 | isa = PBXVariantGroup; 226 | children = ( 227 | 2139409822D32D2400B23738 /* Base */, 228 | ); 229 | name = LaunchScreen.storyboard; 230 | sourceTree = ""; 231 | }; 232 | /* End PBXVariantGroup section */ 233 | 234 | /* Begin XCBuildConfiguration section */ 235 | 2139409B22D32D2400B23738 /* Debug */ = { 236 | isa = XCBuildConfiguration; 237 | buildSettings = { 238 | ALWAYS_SEARCH_USER_PATHS = NO; 239 | CLANG_ANALYZER_NONNULL = YES; 240 | CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; 241 | CLANG_CXX_LANGUAGE_STANDARD = "gnu++14"; 242 | CLANG_CXX_LIBRARY = "libc++"; 243 | CLANG_ENABLE_MODULES = YES; 244 | CLANG_ENABLE_OBJC_ARC = YES; 245 | CLANG_ENABLE_OBJC_WEAK = YES; 246 | CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; 247 | CLANG_WARN_BOOL_CONVERSION = YES; 248 | CLANG_WARN_COMMA = YES; 249 | CLANG_WARN_CONSTANT_CONVERSION = YES; 250 | CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; 251 | CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; 252 | CLANG_WARN_DOCUMENTATION_COMMENTS = YES; 253 | CLANG_WARN_EMPTY_BODY = YES; 254 | CLANG_WARN_ENUM_CONVERSION = YES; 255 | CLANG_WARN_INFINITE_RECURSION = YES; 256 | CLANG_WARN_INT_CONVERSION = YES; 257 | CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; 258 | CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; 259 | CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; 260 | CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; 261 | CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; 262 | CLANG_WARN_STRICT_PROTOTYPES = YES; 263 | CLANG_WARN_SUSPICIOUS_MOVE = YES; 264 | CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; 265 | CLANG_WARN_UNREACHABLE_CODE = YES; 266 | CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; 267 | CODE_SIGN_IDENTITY = "iPhone Developer"; 268 | COPY_PHASE_STRIP = NO; 269 | DEBUG_INFORMATION_FORMAT = dwarf; 270 | ENABLE_STRICT_OBJC_MSGSEND = YES; 271 | ENABLE_TESTABILITY = YES; 272 | GCC_C_LANGUAGE_STANDARD = gnu11; 273 | GCC_DYNAMIC_NO_PIC = NO; 274 | GCC_NO_COMMON_BLOCKS = YES; 275 | GCC_OPTIMIZATION_LEVEL = 0; 276 | GCC_PREPROCESSOR_DEFINITIONS = ( 277 | "DEBUG=1", 278 | "$(inherited)", 279 | ); 280 | GCC_WARN_64_TO_32_BIT_CONVERSION = YES; 281 | GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; 282 | GCC_WARN_UNDECLARED_SELECTOR = YES; 283 | GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; 284 | GCC_WARN_UNUSED_FUNCTION = YES; 285 | GCC_WARN_UNUSED_VARIABLE = YES; 286 | IPHONEOS_DEPLOYMENT_TARGET = 12.0; 287 | MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; 288 | MTL_FAST_MATH = YES; 289 | ONLY_ACTIVE_ARCH = YES; 290 | SDKROOT = iphoneos; 291 | SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG; 292 | SWIFT_OPTIMIZATION_LEVEL = "-Onone"; 293 | }; 294 | name = Debug; 295 | }; 296 | 2139409C22D32D2400B23738 /* Release */ = { 297 | isa = XCBuildConfiguration; 298 | buildSettings = { 299 | ALWAYS_SEARCH_USER_PATHS = NO; 300 | CLANG_ANALYZER_NONNULL = YES; 301 | CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; 302 | CLANG_CXX_LANGUAGE_STANDARD = "gnu++14"; 303 | CLANG_CXX_LIBRARY = "libc++"; 304 | CLANG_ENABLE_MODULES = YES; 305 | CLANG_ENABLE_OBJC_ARC = YES; 306 | CLANG_ENABLE_OBJC_WEAK = YES; 307 | CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; 308 | CLANG_WARN_BOOL_CONVERSION = YES; 309 | CLANG_WARN_COMMA = YES; 310 | CLANG_WARN_CONSTANT_CONVERSION = YES; 311 | CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; 312 | CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; 313 | CLANG_WARN_DOCUMENTATION_COMMENTS = YES; 314 | CLANG_WARN_EMPTY_BODY = YES; 315 | CLANG_WARN_ENUM_CONVERSION = YES; 316 | CLANG_WARN_INFINITE_RECURSION = YES; 317 | CLANG_WARN_INT_CONVERSION = YES; 318 | CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; 319 | CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; 320 | CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; 321 | CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; 322 | CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; 323 | CLANG_WARN_STRICT_PROTOTYPES = YES; 324 | CLANG_WARN_SUSPICIOUS_MOVE = YES; 325 | CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; 326 | CLANG_WARN_UNREACHABLE_CODE = YES; 327 | CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; 328 | CODE_SIGN_IDENTITY = "iPhone Developer"; 329 | COPY_PHASE_STRIP = NO; 330 | DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; 331 | ENABLE_NS_ASSERTIONS = NO; 332 | ENABLE_STRICT_OBJC_MSGSEND = YES; 333 | GCC_C_LANGUAGE_STANDARD = gnu11; 334 | GCC_NO_COMMON_BLOCKS = YES; 335 | GCC_WARN_64_TO_32_BIT_CONVERSION = YES; 336 | GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; 337 | GCC_WARN_UNDECLARED_SELECTOR = YES; 338 | GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; 339 | GCC_WARN_UNUSED_FUNCTION = YES; 340 | GCC_WARN_UNUSED_VARIABLE = YES; 341 | IPHONEOS_DEPLOYMENT_TARGET = 12.0; 342 | MTL_ENABLE_DEBUG_INFO = NO; 343 | MTL_FAST_MATH = YES; 344 | SDKROOT = iphoneos; 345 | SWIFT_COMPILATION_MODE = wholemodule; 346 | SWIFT_OPTIMIZATION_LEVEL = "-O"; 347 | VALIDATE_PRODUCT = YES; 348 | }; 349 | name = Release; 350 | }; 351 | 2139409E22D32D2400B23738 /* Debug */ = { 352 | isa = XCBuildConfiguration; 353 | buildSettings = { 354 | ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; 355 | CODE_SIGN_STYLE = Automatic; 356 | DEVELOPMENT_TEAM = K22D3CEE8V; 357 | INFOPLIST_FILE = BackgroundRemovalWithCoreMLSample/Info.plist; 358 | LD_RUNPATH_SEARCH_PATHS = ( 359 | "$(inherited)", 360 | "@executable_path/Frameworks", 361 | ); 362 | PRODUCT_BUNDLE_IDENTIFIER = com.greenboss.first; 363 | PRODUCT_NAME = "$(TARGET_NAME)"; 364 | SWIFT_VERSION = 5.0; 365 | TARGETED_DEVICE_FAMILY = "1,2"; 366 | }; 367 | name = Debug; 368 | }; 369 | 2139409F22D32D2400B23738 /* Release */ = { 370 | isa = XCBuildConfiguration; 371 | buildSettings = { 372 | ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; 373 | CODE_SIGN_STYLE = Automatic; 374 | DEVELOPMENT_TEAM = K22D3CEE8V; 375 | INFOPLIST_FILE = BackgroundRemovalWithCoreMLSample/Info.plist; 376 | LD_RUNPATH_SEARCH_PATHS = ( 377 | "$(inherited)", 378 | "@executable_path/Frameworks", 379 | ); 380 | PRODUCT_BUNDLE_IDENTIFIER = com.greenboss.first; 381 | PRODUCT_NAME = "$(TARGET_NAME)"; 382 | SWIFT_VERSION = 5.0; 383 | TARGETED_DEVICE_FAMILY = "1,2"; 384 | }; 385 | name = Release; 386 | }; 387 | /* End XCBuildConfiguration section */ 388 | 389 | /* Begin XCConfigurationList section */ 390 | 2139408622D32D2300B23738 /* Build configuration list for PBXProject "BackgroundRemovalWithCoreMLSample" */ = { 391 | isa = XCConfigurationList; 392 | buildConfigurations = ( 393 | 2139409B22D32D2400B23738 /* Debug */, 394 | 2139409C22D32D2400B23738 /* Release */, 395 | ); 396 | defaultConfigurationIsVisible = 0; 397 | defaultConfigurationName = Release; 398 | }; 399 | 2139409D22D32D2400B23738 /* Build configuration list for PBXNativeTarget "BackgroundRemovalWithCoreMLSample" */ = { 400 | isa = XCConfigurationList; 401 | buildConfigurations = ( 402 | 2139409E22D32D2400B23738 /* Debug */, 403 | 2139409F22D32D2400B23738 /* Release */, 404 | ); 405 | defaultConfigurationIsVisible = 0; 406 | defaultConfigurationName = Release; 407 | }; 408 | /* End XCConfigurationList section */ 409 | }; 410 | rootObject = 2139408322D32D2300B23738 /* Project object */; 411 | } 412 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample.xcodeproj/project.xcworkspace/contents.xcworkspacedata: -------------------------------------------------------------------------------- 1 | 2 | 4 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | IDEDidComputeMac32BitWarning 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/AppDelegate.swift: -------------------------------------------------------------------------------- 1 | /* 2 | MIT License 3 | 4 | Copyright (c) 2019 tbchen 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. 23 | */ 24 | 25 | import UIKit 26 | 27 | @UIApplicationMain 28 | class AppDelegate: UIResponder, UIApplicationDelegate { 29 | 30 | var window: UIWindow? 31 | 32 | func application(_ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]?) -> Bool { 33 | return true 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Assets.xcassets/AppIcon.appiconset/Contents.json: -------------------------------------------------------------------------------- 1 | { 2 | "images" : [ 3 | { 4 | "idiom" : "iphone", 5 | "size" : "20x20", 6 | "scale" : "2x" 7 | }, 8 | { 9 | "idiom" : "iphone", 10 | "size" : "20x20", 11 | "scale" : "3x" 12 | }, 13 | { 14 | "idiom" : "iphone", 15 | "size" : "29x29", 16 | "scale" : "2x" 17 | }, 18 | { 19 | "idiom" : "iphone", 20 | "size" : "29x29", 21 | "scale" : "3x" 22 | }, 23 | { 24 | "idiom" : "iphone", 25 | "size" : "40x40", 26 | "scale" : "2x" 27 | }, 28 | { 29 | "idiom" : "iphone", 30 | "size" : "40x40", 31 | "scale" : "3x" 32 | }, 33 | { 34 | "idiom" : "iphone", 35 | "size" : "60x60", 36 | "scale" : "2x" 37 | }, 38 | { 39 | "idiom" : "iphone", 40 | "size" : "60x60", 41 | "scale" : "3x" 42 | }, 43 | { 44 | "idiom" : "ipad", 45 | "size" : "20x20", 46 | "scale" : "1x" 47 | }, 48 | { 49 | "idiom" : "ipad", 50 | "size" : "20x20", 51 | "scale" : "2x" 52 | }, 53 | { 54 | "idiom" : "ipad", 55 | "size" : "29x29", 56 | "scale" : "1x" 57 | }, 58 | { 59 | "idiom" : "ipad", 60 | "size" : "29x29", 61 | "scale" : "2x" 62 | }, 63 | { 64 | "idiom" : "ipad", 65 | "size" : "40x40", 66 | "scale" : "1x" 67 | }, 68 | { 69 | "idiom" : "ipad", 70 | "size" : "40x40", 71 | "scale" : "2x" 72 | }, 73 | { 74 | "idiom" : "ipad", 75 | "size" : "76x76", 76 | "scale" : "1x" 77 | }, 78 | { 79 | "idiom" : "ipad", 80 | "size" : "76x76", 81 | "scale" : "2x" 82 | }, 83 | { 84 | "idiom" : "ipad", 85 | "size" : "83.5x83.5", 86 | "scale" : "2x" 87 | }, 88 | { 89 | "idiom" : "ios-marketing", 90 | "size" : "1024x1024", 91 | "scale" : "1x" 92 | } 93 | ], 94 | "info" : { 95 | "version" : 1, 96 | "author" : "xcode" 97 | } 98 | } -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Assets.xcassets/Contents.json: -------------------------------------------------------------------------------- 1 | { 2 | "info" : { 3 | "version" : 1, 4 | "author" : "xcode" 5 | } 6 | } -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Assets.xcassets/cat.imageset/5122257268_42d3a32729_z.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tbchen/BackgroundRemovalWithCoreMLSample/64713f5dd67ba43499bc968da899d336b189cb1e/BackgroundRemovalWithCoreMLSample/Assets.xcassets/cat.imageset/5122257268_42d3a32729_z.jpg -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Assets.xcassets/cat.imageset/Contents.json: -------------------------------------------------------------------------------- 1 | { 2 | "images" : [ 3 | { 4 | "idiom" : "universal", 5 | "filename" : "5122257268_42d3a32729_z.jpg", 6 | "scale" : "1x" 7 | }, 8 | { 9 | "idiom" : "universal", 10 | "scale" : "2x" 11 | }, 12 | { 13 | "idiom" : "universal", 14 | "scale" : "3x" 15 | } 16 | ], 17 | "info" : { 18 | "version" : 1, 19 | "author" : "xcode" 20 | } 21 | } -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Base.lproj/LaunchScreen.storyboard: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Base.lproj/Main.storyboard: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/Array+Extensions.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Swift 24 | 25 | extension Array where Element: Comparable { 26 | /** 27 | Returns the index and value of the largest element in the array. 28 | 29 | - Note: This method is slow. For faster results, use the standalone 30 | version of argmax() instead. 31 | */ 32 | public func argmax() -> (Int, Element) { 33 | precondition(self.count > 0) 34 | var maxIndex = 0 35 | var maxValue = self[0] 36 | for i in 1.. maxValue { 37 | maxValue = self[i] 38 | maxIndex = i 39 | } 40 | return (maxIndex, maxValue) 41 | } 42 | 43 | /** 44 | Returns the indices of the array's elements in sorted order. 45 | */ 46 | public func argsort(by areInIncreasingOrder: (Element, Element) -> Bool) -> [Array.Index] { 47 | return self.indices.sorted { areInIncreasingOrder(self[$0], self[$1]) } 48 | } 49 | 50 | /** 51 | Returns a new array containing the elements at the specified indices. 52 | */ 53 | public func gather(indices: [Array.Index]) -> [Element] { 54 | return indices.map { self[$0] } 55 | } 56 | } 57 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/CGImage+CVPixelBuffer.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import CoreGraphics 24 | import CoreImage 25 | import VideoToolbox 26 | 27 | extension CGImage { 28 | /** 29 | Resizes the image to width x height and converts it to an RGB CVPixelBuffer. 30 | */ 31 | public func pixelBuffer(width: Int, height: Int, 32 | orientation: CGImagePropertyOrientation) -> CVPixelBuffer? { 33 | return pixelBuffer(width: width, height: height, 34 | pixelFormatType: kCVPixelFormatType_32ARGB, 35 | colorSpace: CGColorSpaceCreateDeviceRGB(), 36 | alphaInfo: .noneSkipFirst, 37 | orientation: orientation) 38 | } 39 | 40 | /** 41 | Resizes the image to width x height and converts it to a grayscale CVPixelBuffer. 42 | */ 43 | public func pixelBufferGray(width: Int, height: Int, 44 | orientation: CGImagePropertyOrientation) -> CVPixelBuffer? { 45 | return pixelBuffer(width: width, height: height, 46 | pixelFormatType: kCVPixelFormatType_OneComponent8, 47 | colorSpace: CGColorSpaceCreateDeviceGray(), 48 | alphaInfo: .none, 49 | orientation: orientation) 50 | } 51 | 52 | func pixelBuffer(width: Int, height: Int, pixelFormatType: OSType, 53 | colorSpace: CGColorSpace, alphaInfo: CGImageAlphaInfo, 54 | orientation: CGImagePropertyOrientation) -> CVPixelBuffer? { 55 | 56 | // TODO: If the orientation is not .up, then rotate the CGImage. 57 | // See also: https://stackoverflow.com/a/40438893/ 58 | assert(orientation == .up) 59 | 60 | var maybePixelBuffer: CVPixelBuffer? 61 | let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, 62 | kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] 63 | let status = CVPixelBufferCreate(kCFAllocatorDefault, 64 | width, 65 | height, 66 | pixelFormatType, 67 | attrs as CFDictionary, 68 | &maybePixelBuffer) 69 | 70 | guard status == kCVReturnSuccess, let pixelBuffer = maybePixelBuffer else { 71 | return nil 72 | } 73 | 74 | let flags = CVPixelBufferLockFlags(rawValue: 0) 75 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(pixelBuffer, flags) else { 76 | return nil 77 | } 78 | defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, flags) } 79 | 80 | guard let context = CGContext(data: CVPixelBufferGetBaseAddress(pixelBuffer), 81 | width: width, 82 | height: height, 83 | bitsPerComponent: 8, 84 | bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer), 85 | space: colorSpace, 86 | bitmapInfo: alphaInfo.rawValue) 87 | else { 88 | return nil 89 | } 90 | 91 | context.draw(self, in: CGRect(x: 0, y: 0, width: width, height: height)) 92 | return pixelBuffer 93 | } 94 | } 95 | 96 | extension CGImage { 97 | /** 98 | Creates a new CGImage from a CVPixelBuffer. 99 | 100 | - Note: Not all CVPixelBuffer pixel formats support conversion into a 101 | CGImage-compatible pixel format. 102 | */ 103 | public static func create(pixelBuffer: CVPixelBuffer) -> CGImage? { 104 | var cgImage: CGImage? 105 | VTCreateCGImageFromCVPixelBuffer(pixelBuffer, options: nil, imageOut: &cgImage) 106 | return cgImage 107 | } 108 | 109 | /* 110 | // Alternative implementation: 111 | public static func create(pixelBuffer: CVPixelBuffer) -> CGImage? { 112 | // This method creates a bitmap CGContext using the pixel buffer's memory. 113 | // It currently only handles kCVPixelFormatType_32ARGB images. To support 114 | // other pixel formats too, you'll have to change the bitmapInfo and maybe 115 | // the color space for the CGContext. 116 | 117 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly) else { 118 | return nil 119 | } 120 | defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) } 121 | 122 | if let context = CGContext(data: CVPixelBufferGetBaseAddress(pixelBuffer), 123 | width: CVPixelBufferGetWidth(pixelBuffer), 124 | height: CVPixelBufferGetHeight(pixelBuffer), 125 | bitsPerComponent: 8, 126 | bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer), 127 | space: CGColorSpaceCreateDeviceRGB(), 128 | bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue), 129 | let cgImage = context.makeImage() { 130 | return cgImage 131 | } else { 132 | return nil 133 | } 134 | } 135 | */ 136 | 137 | /** 138 | Creates a new CGImage from a CVPixelBuffer, using Core Image. 139 | */ 140 | public static func create(pixelBuffer: CVPixelBuffer, context: CIContext) -> CGImage? { 141 | let ciImage = CIImage(cvPixelBuffer: pixelBuffer) 142 | let rect = CGRect(x: 0, y: 0, width: CVPixelBufferGetWidth(pixelBuffer), 143 | height: CVPixelBufferGetHeight(pixelBuffer)) 144 | return context.createCGImage(ciImage, from: rect) 145 | } 146 | } 147 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/CGImage+RawBytes.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import CoreGraphics 24 | 25 | extension CGImage { 26 | /** 27 | Converts the image into an array of RGBA bytes. 28 | */ 29 | @nonobjc public func toByteArrayRGBA() -> [UInt8] { 30 | var bytes = [UInt8](repeating: 0, count: width * height * 4) 31 | bytes.withUnsafeMutableBytes { ptr in 32 | if let colorSpace = colorSpace, 33 | let context = CGContext( 34 | data: ptr.baseAddress, 35 | width: width, 36 | height: height, 37 | bitsPerComponent: bitsPerComponent, 38 | bytesPerRow: bytesPerRow, 39 | space: colorSpace, 40 | bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) { 41 | let rect = CGRect(x: 0, y: 0, width: width, height: height) 42 | context.draw(self, in: rect) 43 | } 44 | } 45 | return bytes 46 | } 47 | 48 | /** 49 | Creates a new CGImage from an array of RGBA bytes. 50 | */ 51 | @nonobjc public class func fromByteArrayRGBA(_ bytes: [UInt8], 52 | width: Int, 53 | height: Int) -> CGImage? { 54 | return fromByteArray(bytes, width: width, height: height, 55 | bytesPerRow: width * 4, 56 | colorSpace: CGColorSpaceCreateDeviceRGB(), 57 | alphaInfo: .premultipliedLast) 58 | } 59 | 60 | /** 61 | Creates a new CGImage from an array of grayscale bytes. 62 | */ 63 | @nonobjc public class func fromByteArrayGray(_ bytes: [UInt8], 64 | width: Int, 65 | height: Int) -> CGImage? { 66 | return fromByteArray(bytes, width: width, height: height, 67 | bytesPerRow: width, 68 | colorSpace: CGColorSpaceCreateDeviceGray(), 69 | alphaInfo: .none) 70 | } 71 | 72 | @nonobjc class func fromByteArray(_ bytes: [UInt8], 73 | width: Int, 74 | height: Int, 75 | bytesPerRow: Int, 76 | colorSpace: CGColorSpace, 77 | alphaInfo: CGImageAlphaInfo) -> CGImage? { 78 | return bytes.withUnsafeBytes { ptr in 79 | let context = CGContext(data: UnsafeMutableRawPointer(mutating: ptr.baseAddress!), 80 | width: width, 81 | height: height, 82 | bitsPerComponent: 8, 83 | bytesPerRow: bytesPerRow, 84 | space: colorSpace, 85 | bitmapInfo: alphaInfo.rawValue) 86 | return context?.makeImage() 87 | } 88 | } 89 | } 90 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/CGImagePropertyOrientation.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | #if canImport(UIKit) 24 | 25 | import UIKit 26 | 27 | public extension CGImagePropertyOrientation { 28 | init(_ orientation: UIImage.Orientation) { 29 | switch orientation { 30 | case .up: self = .up 31 | case .upMirrored: self = .upMirrored 32 | case .down: self = .down 33 | case .downMirrored: self = .downMirrored 34 | case .left: self = .left 35 | case .leftMirrored: self = .leftMirrored 36 | case .right: self = .right 37 | case .rightMirrored: self = .rightMirrored 38 | @unknown default: self = .up 39 | } 40 | } 41 | } 42 | 43 | #if !os(tvOS) 44 | 45 | public extension CGImagePropertyOrientation { 46 | init(_ orientation: UIDeviceOrientation) { 47 | switch orientation { 48 | case .portraitUpsideDown: self = .left 49 | case .landscapeLeft: self = .up 50 | case .landscapeRight: self = .down 51 | default: self = .right 52 | } 53 | } 54 | } 55 | 56 | #endif 57 | 58 | extension UIImage.Orientation { 59 | init(_ cgOrientation: UIImage.Orientation) { 60 | switch cgOrientation { 61 | case .up: self = .up 62 | case .upMirrored: self = .upMirrored 63 | case .down: self = .down 64 | case .downMirrored: self = .downMirrored 65 | case .left: self = .left 66 | case .leftMirrored: self = .leftMirrored 67 | case .right: self = .right 68 | case .rightMirrored: self = .rightMirrored 69 | @unknown default: self = .up 70 | } 71 | } 72 | } 73 | 74 | #endif 75 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/CVPixelBuffer+Helpers.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Foundation 24 | import Accelerate 25 | import CoreImage 26 | 27 | /** 28 | Creates a RGB pixel buffer of the specified width and height. 29 | */ 30 | public func createPixelBuffer(width: Int, height: Int) -> CVPixelBuffer? { 31 | var pixelBuffer: CVPixelBuffer? 32 | let status = CVPixelBufferCreate(nil, width, height, 33 | kCVPixelFormatType_32BGRA, nil, 34 | &pixelBuffer) 35 | if status != kCVReturnSuccess { 36 | print("Error: could not create pixel buffer", status) 37 | return nil 38 | } 39 | return pixelBuffer 40 | } 41 | 42 | /** 43 | First crops the pixel buffer, then resizes it. 44 | 45 | - Note: The new CVPixelBuffer is not backed by an IOSurface and therefore 46 | cannot be turned into a Metal texture. 47 | */ 48 | public func resizePixelBuffer(_ srcPixelBuffer: CVPixelBuffer, 49 | cropX: Int, 50 | cropY: Int, 51 | cropWidth: Int, 52 | cropHeight: Int, 53 | scaleWidth: Int, 54 | scaleHeight: Int) -> CVPixelBuffer? { 55 | let flags = CVPixelBufferLockFlags(rawValue: 0) 56 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(srcPixelBuffer, flags) else { 57 | return nil 58 | } 59 | defer { CVPixelBufferUnlockBaseAddress(srcPixelBuffer, flags) } 60 | 61 | guard let srcData = CVPixelBufferGetBaseAddress(srcPixelBuffer) else { 62 | print("Error: could not get pixel buffer base address") 63 | return nil 64 | } 65 | let srcBytesPerRow = CVPixelBufferGetBytesPerRow(srcPixelBuffer) 66 | let offset = cropY*srcBytesPerRow + cropX*4 67 | var srcBuffer = vImage_Buffer(data: srcData.advanced(by: offset), 68 | height: vImagePixelCount(cropHeight), 69 | width: vImagePixelCount(cropWidth), 70 | rowBytes: srcBytesPerRow) 71 | 72 | let destBytesPerRow = scaleWidth*4 73 | guard let destData = malloc(scaleHeight*destBytesPerRow) else { 74 | print("Error: out of memory") 75 | return nil 76 | } 77 | var destBuffer = vImage_Buffer(data: destData, 78 | height: vImagePixelCount(scaleHeight), 79 | width: vImagePixelCount(scaleWidth), 80 | rowBytes: destBytesPerRow) 81 | 82 | let error = vImageScale_ARGB8888(&srcBuffer, &destBuffer, nil, vImage_Flags(0)) 83 | if error != kvImageNoError { 84 | print("Error:", error) 85 | free(destData) 86 | return nil 87 | } 88 | 89 | let releaseCallback: CVPixelBufferReleaseBytesCallback = { _, ptr in 90 | if let ptr = ptr { 91 | free(UnsafeMutableRawPointer(mutating: ptr)) 92 | } 93 | } 94 | 95 | let pixelFormat = CVPixelBufferGetPixelFormatType(srcPixelBuffer) 96 | var dstPixelBuffer: CVPixelBuffer? 97 | let status = CVPixelBufferCreateWithBytes(nil, scaleWidth, scaleHeight, 98 | pixelFormat, destData, 99 | destBytesPerRow, releaseCallback, 100 | nil, nil, &dstPixelBuffer) 101 | if status != kCVReturnSuccess { 102 | print("Error: could not create new pixel buffer") 103 | free(destData) 104 | return nil 105 | } 106 | return dstPixelBuffer 107 | } 108 | 109 | /** 110 | Resizes a CVPixelBuffer to a new width and height. 111 | 112 | - Note: The new CVPixelBuffer is not backed by an IOSurface and therefore 113 | cannot be turned into a Metal texture. 114 | */ 115 | public func resizePixelBuffer(_ pixelBuffer: CVPixelBuffer, 116 | width: Int, height: Int) -> CVPixelBuffer? { 117 | return resizePixelBuffer(pixelBuffer, cropX: 0, cropY: 0, 118 | cropWidth: CVPixelBufferGetWidth(pixelBuffer), 119 | cropHeight: CVPixelBufferGetHeight(pixelBuffer), 120 | scaleWidth: width, scaleHeight: height) 121 | } 122 | 123 | /** 124 | Resizes a CVPixelBuffer to a new width and height. 125 | */ 126 | public func resizePixelBuffer(_ pixelBuffer: CVPixelBuffer, 127 | width: Int, height: Int, 128 | output: CVPixelBuffer, context: CIContext) { 129 | let ciImage = CIImage(cvPixelBuffer: pixelBuffer) 130 | let sx = CGFloat(width) / CGFloat(CVPixelBufferGetWidth(pixelBuffer)) 131 | let sy = CGFloat(height) / CGFloat(CVPixelBufferGetHeight(pixelBuffer)) 132 | let scaleTransform = CGAffineTransform(scaleX: sx, y: sy) 133 | let scaledImage = ciImage.transformed(by: scaleTransform) 134 | context.render(scaledImage, to: output) 135 | } 136 | 137 | /** 138 | Rotates CVPixelBuffer by the provided factor of 90 counterclock-wise. 139 | 140 | - Note: The new CVPixelBuffer is not backed by an IOSurface and therefore 141 | cannot be turned into a Metal texture. 142 | */ 143 | public func rotate90PixelBuffer(_ srcPixelBuffer: CVPixelBuffer, factor: UInt8) -> CVPixelBuffer? { 144 | let flags = CVPixelBufferLockFlags(rawValue: 0) 145 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(srcPixelBuffer, flags) else { 146 | return nil 147 | } 148 | defer { CVPixelBufferUnlockBaseAddress(srcPixelBuffer, flags) } 149 | 150 | guard let srcData = CVPixelBufferGetBaseAddress(srcPixelBuffer) else { 151 | print("Error: could not get pixel buffer base address") 152 | return nil 153 | } 154 | let sourceWidth = CVPixelBufferGetWidth(srcPixelBuffer) 155 | let sourceHeight = CVPixelBufferGetHeight(srcPixelBuffer) 156 | var destWidth = sourceHeight 157 | var destHeight = sourceWidth 158 | var color = UInt8(0) 159 | 160 | if factor % 2 == 0 { 161 | destWidth = sourceWidth 162 | destHeight = sourceHeight 163 | } 164 | 165 | let srcBytesPerRow = CVPixelBufferGetBytesPerRow(srcPixelBuffer) 166 | var srcBuffer = vImage_Buffer(data: srcData, 167 | height: vImagePixelCount(sourceHeight), 168 | width: vImagePixelCount(sourceWidth), 169 | rowBytes: srcBytesPerRow) 170 | 171 | let destBytesPerRow = destWidth*4 172 | guard let destData = malloc(destHeight*destBytesPerRow) else { 173 | print("Error: out of memory") 174 | return nil 175 | } 176 | var destBuffer = vImage_Buffer(data: destData, 177 | height: vImagePixelCount(destHeight), 178 | width: vImagePixelCount(destWidth), 179 | rowBytes: destBytesPerRow) 180 | 181 | let error = vImageRotate90_ARGB8888(&srcBuffer, &destBuffer, factor, &color, vImage_Flags(0)) 182 | if error != kvImageNoError { 183 | print("Error:", error) 184 | free(destData) 185 | return nil 186 | } 187 | 188 | let releaseCallback: CVPixelBufferReleaseBytesCallback = { _, ptr in 189 | if let ptr = ptr { 190 | free(UnsafeMutableRawPointer(mutating: ptr)) 191 | } 192 | } 193 | 194 | let pixelFormat = CVPixelBufferGetPixelFormatType(srcPixelBuffer) 195 | var dstPixelBuffer: CVPixelBuffer? 196 | let status = CVPixelBufferCreateWithBytes(nil, destWidth, destHeight, 197 | pixelFormat, destData, 198 | destBytesPerRow, releaseCallback, 199 | nil, nil, &dstPixelBuffer) 200 | if status != kCVReturnSuccess { 201 | print("Error: could not create new pixel buffer") 202 | free(destData) 203 | return nil 204 | } 205 | return dstPixelBuffer 206 | } 207 | 208 | public extension CVPixelBuffer { 209 | /** 210 | Copies a CVPixelBuffer to a new CVPixelBuffer that is compatible with Metal. 211 | 212 | - Tip: If CVMetalTextureCacheCreateTextureFromImage is failing, then call 213 | this method first! 214 | */ 215 | func copyToMetalCompatible() -> CVPixelBuffer? { 216 | // Other possible options: 217 | // String(kCVPixelBufferOpenGLCompatibilityKey): true, 218 | // String(kCVPixelBufferIOSurfacePropertiesKey): [ 219 | // "IOSurfaceOpenGLESFBOCompatibility": true, 220 | // "IOSurfaceOpenGLESTextureCompatibility": true, 221 | // "IOSurfaceCoreAnimationCompatibility": true 222 | // ] 223 | let attributes: [String: Any] = [ 224 | String(kCVPixelBufferMetalCompatibilityKey): true, 225 | ] 226 | return deepCopy(withAttributes: attributes) 227 | } 228 | 229 | /** 230 | Copies a CVPixelBuffer to a new CVPixelBuffer. 231 | 232 | This lets you specify new attributes, such as whether the new CVPixelBuffer 233 | must be IOSurface-backed. 234 | 235 | See: https://developer.apple.com/library/archive/qa/qa1781/_index.html 236 | */ 237 | func deepCopy(withAttributes attributes: [String: Any] = [:]) -> CVPixelBuffer? { 238 | let srcPixelBuffer = self 239 | let srcFlags: CVPixelBufferLockFlags = .readOnly 240 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(srcPixelBuffer, srcFlags) else { 241 | return nil 242 | } 243 | defer { CVPixelBufferUnlockBaseAddress(srcPixelBuffer, srcFlags) } 244 | 245 | var combinedAttributes: [String: Any] = [:] 246 | 247 | // Copy attachment attributes. 248 | if let attachments = CVBufferGetAttachments(srcPixelBuffer, .shouldPropagate) as? [String: Any] { 249 | for (key, value) in attachments { 250 | combinedAttributes[key] = value 251 | } 252 | } 253 | 254 | // Add user attributes. 255 | combinedAttributes = combinedAttributes.merging(attributes) { $1 } 256 | 257 | var maybePixelBuffer: CVPixelBuffer? 258 | let status = CVPixelBufferCreate(kCFAllocatorDefault, 259 | CVPixelBufferGetWidth(srcPixelBuffer), 260 | CVPixelBufferGetHeight(srcPixelBuffer), 261 | CVPixelBufferGetPixelFormatType(srcPixelBuffer), 262 | combinedAttributes as CFDictionary, 263 | &maybePixelBuffer) 264 | 265 | guard status == kCVReturnSuccess, let dstPixelBuffer = maybePixelBuffer else { 266 | return nil 267 | } 268 | 269 | let dstFlags = CVPixelBufferLockFlags(rawValue: 0) 270 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(dstPixelBuffer, dstFlags) else { 271 | return nil 272 | } 273 | defer { CVPixelBufferUnlockBaseAddress(dstPixelBuffer, dstFlags) } 274 | 275 | for plane in 0...max(0, CVPixelBufferGetPlaneCount(srcPixelBuffer) - 1) { 276 | if let srcAddr = CVPixelBufferGetBaseAddressOfPlane(srcPixelBuffer, plane), 277 | let dstAddr = CVPixelBufferGetBaseAddressOfPlane(dstPixelBuffer, plane) { 278 | let srcBytesPerRow = CVPixelBufferGetBytesPerRowOfPlane(srcPixelBuffer, plane) 279 | let dstBytesPerRow = CVPixelBufferGetBytesPerRowOfPlane(dstPixelBuffer, plane) 280 | 281 | for h in 0.. MLMultiArray { 34 | let newCount = dimensions.reduce(1, *) 35 | precondition(newCount == count, "Cannot reshape \(shape) to \(dimensions)") 36 | 37 | var newStrides = [Int](repeating: 0, count: dimensions.count) 38 | newStrides[dimensions.count - 1] = 1 39 | for i in stride(from: dimensions.count - 1, to: 0, by: -1) { 40 | newStrides[i - 1] = newStrides[i] * dimensions[i] 41 | } 42 | 43 | let newShape_ = dimensions.map { NSNumber(value: $0) } 44 | let newStrides_ = newStrides.map { NSNumber(value: $0) } 45 | 46 | return try MLMultiArray(dataPointer: self.dataPointer, 47 | shape: newShape_, 48 | dataType: self.dataType, 49 | strides: newStrides_) 50 | } 51 | 52 | /** 53 | Returns a transposed version of this MLMultiArray. 54 | 55 | - Note: This copies the data. 56 | 57 | - TODO: Support .float32 and .int32 types too. 58 | */ 59 | @nonobjc public func transposed(to order: [Int]) throws -> MLMultiArray { 60 | let ndim = order.count 61 | 62 | precondition(dataType == .double) 63 | precondition(ndim == strides.count) 64 | 65 | let newShape = shape.indices.map { shape[order[$0]] } 66 | let newArray = try MLMultiArray(shape: newShape, dataType: self.dataType) 67 | 68 | let srcPtr = UnsafeMutablePointer(OpaquePointer(dataPointer)) 69 | let dstPtr = UnsafeMutablePointer(OpaquePointer(newArray.dataPointer)) 70 | 71 | let srcShape = shape.map { $0.intValue } 72 | let dstStride = newArray.strides.map { $0.intValue } 73 | var idx = [Int](repeating: 0, count: ndim) 74 | 75 | for j in 0.. 0 && idx[i] >= srcShape[i] { 89 | idx[i] = 0 90 | idx[i - 1] += 1 91 | i -= 1 92 | } 93 | } 94 | return newArray 95 | } 96 | } 97 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/MLMultiArray+Image.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Accelerate 24 | import CoreML 25 | 26 | public protocol MultiArrayType: Comparable { 27 | static var multiArrayDataType: MLMultiArrayDataType { get } 28 | static func +(lhs: Self, rhs: Self) -> Self 29 | static func -(lhs: Self, rhs: Self) -> Self 30 | static func *(lhs: Self, rhs: Self) -> Self 31 | static func /(lhs: Self, rhs: Self) -> Self 32 | init(_: Int) 33 | var toUInt8: UInt8 { get } 34 | } 35 | 36 | extension Double: MultiArrayType { 37 | public static var multiArrayDataType: MLMultiArrayDataType { return .double } 38 | public var toUInt8: UInt8 { return UInt8(self) } 39 | } 40 | 41 | extension Float: MultiArrayType { 42 | public static var multiArrayDataType: MLMultiArrayDataType { return .float32 } 43 | public var toUInt8: UInt8 { return UInt8(self) } 44 | } 45 | 46 | extension Int32: MultiArrayType { 47 | public static var multiArrayDataType: MLMultiArrayDataType { return .int32 } 48 | public var toUInt8: UInt8 { return UInt8(self) } 49 | } 50 | 51 | extension MLMultiArray { 52 | /** 53 | Converts the multi-array to a CGImage. 54 | 55 | The multi-array must have at least 2 dimensions for a grayscale image, or 56 | at least 3 dimensions for a color image. 57 | 58 | The default expected shape is (height, width) or (channels, height, width). 59 | However, you can change this using the `axes` parameter. For example, if 60 | the array shape is (1, height, width, channels), use `axes: (3, 1, 2)`. 61 | 62 | If `channel` is not nil, only converts that channel to a grayscale image. 63 | This lets you visualize individual channels from a multi-array with more 64 | than 4 channels. 65 | 66 | Otherwise, converts all channels. In this case, the number of channels in 67 | the multi-array must be 1 for grayscale, 3 for RGB, or 4 for RGBA. 68 | 69 | Use the `min` and `max` parameters to put the values from the array into 70 | the range [0, 255], if not already: 71 | 72 | - `min`: should be the smallest value in the data; this will be mapped to 0. 73 | - `max`: should be the largest value in the data; will be mapped to 255. 74 | 75 | For example, if the range of the data in the multi-array is [-1, 1], use 76 | `min: -1, max: 1`. If the range is already [0, 255], then use the defaults. 77 | */ 78 | public func cgImage(min: Double = 0, 79 | max: Double = 255, 80 | channel: Int? = nil, 81 | axes: (Int, Int, Int)? = nil) -> CGImage? { 82 | switch self.dataType { 83 | case .double: 84 | return _image(min: min, max: max, channel: channel, axes: axes) 85 | case .float32: 86 | return _image(min: Float(min), max: Float(max), channel: channel, axes: axes) 87 | case .int32: 88 | return _image(min: Int32(min), max: Int32(max), channel: channel, axes: axes) 89 | @unknown default: 90 | fatalError("Unsupported data type \(dataType.rawValue)") 91 | } 92 | } 93 | 94 | /** 95 | Helper function that allows us to use generics. The type of `min` and `max` 96 | is also the dataType of the MLMultiArray. 97 | */ 98 | private func _image(min: T, 99 | max: T, 100 | channel: Int?, 101 | axes: (Int, Int, Int)?) -> CGImage? { 102 | if let (b, w, h, c) = toRawBytes(min: min, max: max, channel: channel, axes: axes) { 103 | if c == 1 { 104 | return CGImage.fromByteArrayGray(b, width: w, height: h) 105 | } else { 106 | return CGImage.fromByteArrayRGBA(b, width: w, height: h) 107 | } 108 | } 109 | return nil 110 | } 111 | 112 | /** 113 | Converts the multi-array into an array of RGBA or grayscale pixels. 114 | 115 | - Note: This is not particularly fast, but it is flexible. You can change 116 | the loops to convert the multi-array whichever way you please. 117 | 118 | - Note: The type of `min` and `max` must match the dataType of the 119 | MLMultiArray object. 120 | 121 | - Returns: tuple containing the RGBA bytes, the dimensions of the image, 122 | and the number of channels in the image (1, 3, or 4). 123 | */ 124 | public func toRawBytes(min: T, 125 | max: T, 126 | channel: Int? = nil, 127 | axes: (Int, Int, Int)? = nil) 128 | -> (bytes: [UInt8], width: Int, height: Int, channels: Int)? { 129 | // MLMultiArray with unsupported shape? 130 | if shape.count < 2 { 131 | print("Cannot convert MLMultiArray of shape \(shape) to image") 132 | return nil 133 | } 134 | 135 | // Figure out which dimensions to use for the channels, height, and width. 136 | let channelAxis: Int 137 | let heightAxis: Int 138 | let widthAxis: Int 139 | if let axes = axes { 140 | channelAxis = axes.0 141 | heightAxis = axes.1 142 | widthAxis = axes.2 143 | guard channelAxis >= 0 && channelAxis < shape.count && 144 | heightAxis >= 0 && heightAxis < shape.count && 145 | widthAxis >= 0 && widthAxis < shape.count else { 146 | print("Invalid axes \(axes) for shape \(shape)") 147 | return nil 148 | } 149 | } else { 150 | channelAxis = 0 151 | heightAxis = 1 152 | widthAxis = 2 153 | } 154 | 155 | let height = self.shape[heightAxis].intValue 156 | let width = self.shape[widthAxis].intValue 157 | let yStride = self.strides[heightAxis].intValue 158 | let xStride = self.strides[widthAxis].intValue 159 | 160 | let channels: Int 161 | let cStride: Int 162 | let bytesPerPixel: Int 163 | let channelOffset: Int 164 | 165 | // MLMultiArray with just two dimensions is always grayscale. (We ignore 166 | // the value of channelAxis here.) 167 | if shape.count == 2 { 168 | channels = 1 169 | cStride = 0 170 | bytesPerPixel = 1 171 | channelOffset = 0 172 | 173 | // MLMultiArray with more than two dimensions can be color or grayscale. 174 | } else { 175 | let channelDim = self.shape[channelAxis].intValue 176 | if let channel = channel { 177 | if channel < 0 || channel >= channelDim { 178 | print("Channel must be -1, or between 0 and \(channelDim - 1)") 179 | return nil 180 | } 181 | channels = 1 182 | bytesPerPixel = 1 183 | channelOffset = channel 184 | } else if channelDim == 1 { 185 | channels = 1 186 | bytesPerPixel = 1 187 | channelOffset = 0 188 | } else { 189 | if channelDim != 3 && channelDim != 4 { 190 | print("Expected channel dimension to have 1, 3, or 4 channels, got \(channelDim)") 191 | return nil 192 | } 193 | channels = channelDim 194 | bytesPerPixel = 4 195 | channelOffset = 0 196 | } 197 | cStride = self.strides[channelAxis].intValue 198 | } 199 | 200 | // Allocate storage for the RGBA or grayscale pixels. Set everything to 201 | // 255 so that alpha channel is filled in if only 3 channels. 202 | let count = height * width * bytesPerPixel 203 | var pixels = [UInt8](repeating: 255, count: count) 204 | 205 | // Grab the pointer to MLMultiArray's memory. 206 | var ptr = UnsafeMutablePointer(OpaquePointer(self.dataPointer)) 207 | ptr = ptr.advanced(by: channelOffset * cStride) 208 | 209 | // Loop through all the pixels and all the channels and copy them over. 210 | for c in 0.. CGImage? { 240 | assert(features.dataType == .float32) 241 | assert(features.shape.count == 3) 242 | 243 | let ptr = UnsafeMutablePointer(OpaquePointer(features.dataPointer)) 244 | 245 | let height = features.shape[1].intValue 246 | let width = features.shape[2].intValue 247 | let channelStride = features.strides[0].intValue 248 | let rowStride = features.strides[1].intValue 249 | let srcRowBytes = rowStride * MemoryLayout.stride 250 | 251 | var blueBuffer = vImage_Buffer(data: ptr, 252 | height: vImagePixelCount(height), 253 | width: vImagePixelCount(width), 254 | rowBytes: srcRowBytes) 255 | var greenBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride), 256 | height: vImagePixelCount(height), 257 | width: vImagePixelCount(width), 258 | rowBytes: srcRowBytes) 259 | var redBuffer = vImage_Buffer(data: ptr.advanced(by: channelStride * 2), 260 | height: vImagePixelCount(height), 261 | width: vImagePixelCount(width), 262 | rowBytes: srcRowBytes) 263 | 264 | let destRowBytes = width * 4 265 | var pixels = [UInt8](repeating: 0, count: height * destRowBytes) 266 | var destBuffer = vImage_Buffer(data: &pixels, 267 | height: vImagePixelCount(height), 268 | width: vImagePixelCount(width), 269 | rowBytes: destRowBytes) 270 | 271 | let error = vImageConvert_PlanarFToBGRX8888(&blueBuffer, 272 | &greenBuffer, 273 | &redBuffer, 274 | Pixel_8(255), 275 | &destBuffer, 276 | [max, max, max], 277 | [min, min, min], 278 | vImage_Flags(0)) 279 | if error == kvImageNoError { 280 | return CGImage.fromByteArrayRGBA(pixels, width: width, height: height) 281 | } else { 282 | return nil 283 | } 284 | } 285 | 286 | #if canImport(UIKit) 287 | 288 | import UIKit 289 | 290 | extension MLMultiArray { 291 | public func image(min: Double = 0, 292 | max: Double = 255, 293 | channel: Int? = nil, 294 | axes: (Int, Int, Int)? = nil) -> UIImage? { 295 | let cgImg = cgImage(min: min, max: max, channel: channel, axes: axes) 296 | return cgImg.map { UIImage(cgImage: $0) } 297 | } 298 | } 299 | 300 | public func createUIImage(fromFloatArray features: MLMultiArray, 301 | min: Float = 0, 302 | max: Float = 255) -> UIImage? { 303 | let cgImg = createCGImage(fromFloatArray: features, min: min, max: max) 304 | return cgImg.map { UIImage(cgImage: $0) } 305 | } 306 | 307 | #endif 308 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/Math.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Foundation 24 | import Accelerate 25 | 26 | /** 27 | Returns the index and value of the largest element in the array. 28 | 29 | - Parameters: 30 | - count: If provided, only look at the first `count` elements of the array, 31 | otherwise look at the entire array. 32 | */ 33 | public func argmax(_ array: [Float], count: Int? = nil) -> (Int, Float) { 34 | var maxValue: Float = 0 35 | var maxIndex: vDSP_Length = 0 36 | vDSP_maxvi(array, 1, &maxValue, &maxIndex, vDSP_Length(count ?? array.count)) 37 | return (Int(maxIndex), maxValue) 38 | } 39 | 40 | /** 41 | Returns the index and value of the largest element in the array. 42 | 43 | - Parameters: 44 | - ptr: Pointer to the first element in memory. 45 | - count: How many elements to look at. 46 | - stride: The distance between two elements in memory. 47 | */ 48 | public func argmax(_ ptr: UnsafePointer, count: Int, stride: Int = 1) -> (Int, Float) { 49 | var maxValue: Float = 0 50 | var maxIndex: vDSP_Length = 0 51 | vDSP_maxvi(ptr, vDSP_Stride(stride), &maxValue, &maxIndex, vDSP_Length(count)) 52 | return (Int(maxIndex), maxValue) 53 | } 54 | 55 | /** 56 | Returns the index and value of the largest element in the array. 57 | 58 | - Parameters: 59 | - count: If provided, only look at the first `count` elements of the array, 60 | otherwise look at the entire array. 61 | */ 62 | public func argmax(_ array: [Double], count: Int? = nil) -> (Int, Double) { 63 | var maxValue: Double = 0 64 | var maxIndex: vDSP_Length = 0 65 | vDSP_maxviD(array, 1, &maxValue, &maxIndex, vDSP_Length(count ?? array.count)) 66 | return (Int(maxIndex), maxValue) 67 | } 68 | 69 | /** 70 | Returns the index and value of the largest element in the array. 71 | 72 | - Parameters: 73 | - ptr: Pointer to the first element in memory. 74 | - count: How many elements to look at. 75 | - stride: The distance between two elements in memory. 76 | */ 77 | public func argmax(_ ptr: UnsafePointer, count: Int, stride: Int = 1) -> (Int, Double) { 78 | var maxValue: Double = 0 79 | var maxIndex: vDSP_Length = 0 80 | vDSP_maxviD(ptr, vDSP_Stride(stride), &maxValue, &maxIndex, vDSP_Length(count)) 81 | return (Int(maxIndex), maxValue) 82 | } 83 | 84 | /** Ensures that `x` is in the range `[min, max]`. */ 85 | public func clamp(_ x: T, min: T, max: T) -> T { 86 | if x < min { return min } 87 | if x > max { return max } 88 | return x 89 | } 90 | 91 | /** Logistic sigmoid. */ 92 | public func sigmoid(_ x: Float) -> Float { 93 | return 1 / (1 + exp(-x)) 94 | } 95 | 96 | /** Logistic sigmoid. */ 97 | public func sigmoid(_ x: Double) -> Double { 98 | return 1 / (1 + exp(-x)) 99 | } 100 | 101 | /* In-place logistic sigmoid: x = 1 / (1 + exp(-x)) */ 102 | public func sigmoid(_ x: UnsafeMutablePointer, count: Int) { 103 | vDSP_vneg(x, 1, x, 1, vDSP_Length(count)) 104 | var cnt = Int32(count) 105 | vvexpf(x, x, &cnt) 106 | var y: Float = 1 107 | vDSP_vsadd(x, 1, &y, x, 1, vDSP_Length(count)) 108 | vvrecf(x, x, &cnt) 109 | } 110 | 111 | /* In-place logistic sigmoid: x = 1 / (1 + exp(-x)) */ 112 | public func sigmoid(_ x: UnsafeMutablePointer, count: Int) { 113 | vDSP_vnegD(x, 1, x, 1, vDSP_Length(count)) 114 | var cnt = Int32(count) 115 | vvexp(x, x, &cnt) 116 | var y: Double = 1 117 | vDSP_vsaddD(x, 1, &y, x, 1, vDSP_Length(count)) 118 | vvrec(x, x, &cnt) 119 | } 120 | 121 | /** 122 | Computes the "softmax" function over an array. 123 | 124 | Based on code from https://github.com/nikolaypavlov/MLPNeuralNet/ 125 | 126 | This is what softmax looks like in "pseudocode" (actually using Python 127 | and numpy): 128 | 129 | x -= np.max(x) 130 | exp_scores = np.exp(x) 131 | softmax = exp_scores / np.sum(exp_scores) 132 | 133 | First we shift the values of x so that the highest value in the array is 0. 134 | This ensures numerical stability with the exponents, so they don't blow up. 135 | */ 136 | public func softmax(_ x: [Float]) -> [Float] { 137 | var x = x 138 | let len = vDSP_Length(x.count) 139 | 140 | // Find the maximum value in the input array. 141 | var max: Float = 0 142 | vDSP_maxv(x, 1, &max, len) 143 | 144 | // Subtract the maximum from all the elements in the array. 145 | // Now the highest value in the array is 0. 146 | max = -max 147 | vDSP_vsadd(x, 1, &max, &x, 1, len) 148 | 149 | // Exponentiate all the elements in the array. 150 | var count = Int32(x.count) 151 | vvexpf(&x, x, &count) 152 | 153 | // Compute the sum of all exponentiated values. 154 | var sum: Float = 0 155 | vDSP_sve(x, 1, &sum, len) 156 | 157 | // Divide each element by the sum. This normalizes the array contents 158 | // so that they all add up to 1. 159 | vDSP_vsdiv(x, 1, &sum, &x, 1, len) 160 | 161 | return x 162 | } 163 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/NonMaxSuppression.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Foundation 24 | import Accelerate 25 | 26 | public struct BoundingBox { 27 | /** Index of the predicted class. */ 28 | public let classIndex: Int 29 | 30 | /** Confidence score. */ 31 | public let score: Float 32 | 33 | /** Normalized coordinates between 0 and 1. */ 34 | public let rect: CGRect 35 | 36 | public init(classIndex: Int, score: Float, rect: CGRect) { 37 | self.classIndex = classIndex 38 | self.score = score 39 | self.rect = rect 40 | } 41 | } 42 | 43 | /** 44 | Computes intersection-over-union overlap between two bounding boxes. 45 | */ 46 | public func IOU(_ a: CGRect, _ b: CGRect) -> Float { 47 | let areaA = a.width * a.height 48 | if areaA <= 0 { return 0 } 49 | 50 | let areaB = b.width * b.height 51 | if areaB <= 0 { return 0 } 52 | 53 | let intersectionMinX = max(a.minX, b.minX) 54 | let intersectionMinY = max(a.minY, b.minY) 55 | let intersectionMaxX = min(a.maxX, b.maxX) 56 | let intersectionMaxY = min(a.maxY, b.maxY) 57 | let intersectionArea = max(intersectionMaxY - intersectionMinY, 0) * 58 | max(intersectionMaxX - intersectionMinX, 0) 59 | return Float(intersectionArea / (areaA + areaB - intersectionArea)) 60 | } 61 | 62 | /** 63 | Removes bounding boxes that overlap too much with other boxes that have 64 | a higher score. 65 | */ 66 | public func nonMaxSuppression(boundingBoxes: [BoundingBox], 67 | iouThreshold: Float, 68 | maxBoxes: Int) -> [Int] { 69 | return nonMaxSuppression(boundingBoxes: boundingBoxes, 70 | indices: Array(boundingBoxes.indices), 71 | iouThreshold: iouThreshold, 72 | maxBoxes: maxBoxes) 73 | } 74 | 75 | /** 76 | Removes bounding boxes that overlap too much with other boxes that have 77 | a higher score. 78 | 79 | Based on code from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/non_max_suppression_op.cc 80 | 81 | - Note: This version of NMS ignores the class of the bounding boxes. Since it 82 | selects the bounding boxes in a greedy fashion, if a certain class has many 83 | boxes that are selected, then it is possible none of the boxes of the other 84 | classes get selected. 85 | 86 | - Parameters: 87 | - boundingBoxes: an array of bounding boxes and their scores 88 | - indices: which predictions to look at 89 | - iouThreshold: used to decide whether boxes overlap too much 90 | - maxBoxes: the maximum number of boxes that will be selected 91 | 92 | - Returns: the array indices of the selected bounding boxes 93 | */ 94 | public func nonMaxSuppression(boundingBoxes: [BoundingBox], 95 | indices: [Int], 96 | iouThreshold: Float, 97 | maxBoxes: Int) -> [Int] { 98 | 99 | // Sort the boxes based on their confidence scores, from high to low. 100 | let sortedIndices = indices.sorted { boundingBoxes[$0].score > boundingBoxes[$1].score } 101 | 102 | var selected: [Int] = [] 103 | 104 | // Loop through the bounding boxes, from highest score to lowest score, 105 | // and determine whether or not to keep each box. 106 | for i in 0..= maxBoxes { break } 108 | 109 | var shouldSelect = true 110 | let boxA = boundingBoxes[sortedIndices[i]] 111 | 112 | // Does the current box overlap one of the selected boxes more than the 113 | // given threshold amount? Then it's too similar, so don't keep it. 114 | for j in 0.. iouThreshold { 117 | shouldSelect = false 118 | break 119 | } 120 | } 121 | 122 | // This bounding box did not overlap too much with any previously selected 123 | // bounding box, so we'll keep it. 124 | if shouldSelect { 125 | selected.append(sortedIndices[i]) 126 | } 127 | } 128 | 129 | return selected 130 | } 131 | 132 | /** 133 | Multi-class version of non maximum suppression. 134 | 135 | Where `nonMaxSuppression()` does not look at the class of the predictions at 136 | all, the multi-class version first selects the best bounding boxes for each 137 | class, and then keeps the best ones of those. 138 | 139 | With this method you can usually expect to see at least one bounding box for 140 | each class (unless all the scores for a given class are really low). 141 | 142 | Based on code from: https://github.com/tensorflow/models/blob/master/object_detection/core/post_processing.py 143 | 144 | - Parameters: 145 | - numClasses: the number of classes 146 | - boundingBoxes: an array of bounding boxes and their scores 147 | - scoreThreshold: used to only keep bounding boxes with a high enough score 148 | - iouThreshold: used to decide whether boxes overlap too much 149 | - maxPerClass: the maximum number of boxes that will be selected per class 150 | - maxTotal: maximum number of boxes that will be selected over all classes 151 | 152 | - Returns: the array indices of the selected bounding boxes 153 | */ 154 | public func nonMaxSuppressionMultiClass(numClasses: Int, 155 | boundingBoxes: [BoundingBox], 156 | scoreThreshold: Float, 157 | iouThreshold: Float, 158 | maxPerClass: Int, 159 | maxTotal: Int) -> [Int] { 160 | var selectedBoxes: [Int] = [] 161 | 162 | // Look at all the classes one-by-one. 163 | for c in 0.. scoreThreshold { 173 | filteredBoxes.append(p) 174 | } 175 | } 176 | } 177 | 178 | // Only keep the best bounding boxes for this class. 179 | let nmsBoxes = nonMaxSuppression(boundingBoxes: boundingBoxes, 180 | indices: filteredBoxes, 181 | iouThreshold: iouThreshold, 182 | maxBoxes: maxPerClass) 183 | 184 | // Add the indices of the surviving boxes to the big list. 185 | selectedBoxes.append(contentsOf: nmsBoxes) 186 | } 187 | 188 | // Sort all the surviving boxes by score and only keep the best ones. 189 | let sortedBoxes = selectedBoxes.sorted { boundingBoxes[$0].score > boundingBoxes[$1].score } 190 | return Array(sortedBoxes.prefix(maxTotal)) 191 | } 192 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/Predictions.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | import Vision 24 | 25 | /** 26 | Returns the top `k` predictions from Core ML classification results as an 27 | array of `(String, Double)` pairs. 28 | */ 29 | public func top(_ k: Int, _ prob: [String: Double]) -> [(String, Double)] { 30 | return Array(prob.map { x in (x.key, x.value) } 31 | .sorted(by: { a, b -> Bool in a.1 > b.1 }) 32 | .prefix(through: min(k, prob.count) - 1)) 33 | } 34 | 35 | /** 36 | Returns the top `k` predictions from Vision classification results as an 37 | array of `(String, Double)` pairs. 38 | */ 39 | public func top(_ k: Int, _ observations: [VNClassificationObservation]) -> [(String, Double)] { 40 | // The Vision observations are sorted by confidence already. 41 | return observations.prefix(through: min(k, observations.count) - 1) 42 | .map { ($0.identifier, Double($0.confidence)) } 43 | } 44 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/UIImage+CVPixelBuffer.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | #if canImport(UIKit) 24 | 25 | import UIKit 26 | import VideoToolbox 27 | 28 | extension UIImage { 29 | /** 30 | Resizes the image to width x height and converts it to an RGB CVPixelBuffer. 31 | */ 32 | public func pixelBuffer(width: Int, height: Int) -> CVPixelBuffer? { 33 | return pixelBuffer(width: width, height: height, 34 | pixelFormatType: kCVPixelFormatType_32ARGB, 35 | colorSpace: CGColorSpaceCreateDeviceRGB(), 36 | alphaInfo: .noneSkipFirst) 37 | } 38 | 39 | /** 40 | Resizes the image to width x height and converts it to a grayscale CVPixelBuffer. 41 | */ 42 | public func pixelBufferGray(width: Int, height: Int) -> CVPixelBuffer? { 43 | return pixelBuffer(width: width, height: height, 44 | pixelFormatType: kCVPixelFormatType_OneComponent8, 45 | colorSpace: CGColorSpaceCreateDeviceGray(), 46 | alphaInfo: .none) 47 | } 48 | 49 | func pixelBuffer(width: Int, height: Int, pixelFormatType: OSType, 50 | colorSpace: CGColorSpace, alphaInfo: CGImageAlphaInfo) -> CVPixelBuffer? { 51 | var maybePixelBuffer: CVPixelBuffer? 52 | let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, 53 | kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] 54 | let status = CVPixelBufferCreate(kCFAllocatorDefault, 55 | width, 56 | height, 57 | pixelFormatType, 58 | attrs as CFDictionary, 59 | &maybePixelBuffer) 60 | 61 | guard status == kCVReturnSuccess, let pixelBuffer = maybePixelBuffer else { 62 | return nil 63 | } 64 | 65 | let flags = CVPixelBufferLockFlags(rawValue: 0) 66 | guard kCVReturnSuccess == CVPixelBufferLockBaseAddress(pixelBuffer, flags) else { 67 | return nil 68 | } 69 | defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, flags) } 70 | 71 | guard let context = CGContext(data: CVPixelBufferGetBaseAddress(pixelBuffer), 72 | width: width, 73 | height: height, 74 | bitsPerComponent: 8, 75 | bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer), 76 | space: colorSpace, 77 | bitmapInfo: alphaInfo.rawValue) 78 | else { 79 | return nil 80 | } 81 | 82 | UIGraphicsPushContext(context) 83 | context.translateBy(x: 0, y: CGFloat(height)) 84 | context.scaleBy(x: 1, y: -1) 85 | self.draw(in: CGRect(x: 0, y: 0, width: width, height: height)) 86 | UIGraphicsPopContext() 87 | 88 | return pixelBuffer 89 | } 90 | } 91 | 92 | extension UIImage { 93 | /** 94 | Creates a new UIImage from a CVPixelBuffer. 95 | 96 | - Note: Not all CVPixelBuffer pixel formats support conversion into a 97 | CGImage-compatible pixel format. 98 | */ 99 | public convenience init?(pixelBuffer: CVPixelBuffer) { 100 | if let cgImage = CGImage.create(pixelBuffer: pixelBuffer) { 101 | self.init(cgImage: cgImage) 102 | } else { 103 | return nil 104 | } 105 | } 106 | 107 | /* 108 | // Alternative implementation: 109 | public convenience init?(pixelBuffer: CVPixelBuffer) { 110 | // This converts the image to a CIImage first and then to a UIImage. 111 | // Does not appear to work on the simulator but is OK on the device. 112 | self.init(ciImage: CIImage(cvPixelBuffer: pixelBuffer)) 113 | } 114 | */ 115 | 116 | /** 117 | Creates a new UIImage from a CVPixelBuffer, using a Core Image context. 118 | */ 119 | public convenience init?(pixelBuffer: CVPixelBuffer, context: CIContext) { 120 | if let cgImage = CGImage.create(pixelBuffer: pixelBuffer, context: context) { 121 | self.init(cgImage: cgImage) 122 | } else { 123 | return nil 124 | } 125 | } 126 | } 127 | 128 | #endif 129 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/UIImage+Extensions.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | #if canImport(UIKit) 24 | 25 | import UIKit 26 | 27 | extension UIImage { 28 | /** 29 | Resizes the image. 30 | 31 | - Parameters: 32 | - scale: If this is 1, `newSize` is the size in pixels. 33 | */ 34 | @nonobjc public func resized(to newSize: CGSize, scale: CGFloat = 1) -> UIImage { 35 | let format = UIGraphicsImageRendererFormat.default() 36 | format.scale = scale 37 | let renderer = UIGraphicsImageRenderer(size: newSize, format: format) 38 | let image = renderer.image { _ in 39 | draw(in: CGRect(origin: .zero, size: newSize)) 40 | } 41 | return image 42 | } 43 | } 44 | 45 | #endif 46 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/CoreMLHelpers/UIImage+RawBytes.swift: -------------------------------------------------------------------------------- 1 | /* 2 | Copyright (c) 2017-2019 M.I. Hollemans 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy 5 | of this software and associated documentation files (the "Software"), to 6 | deal in the Software without restriction, including without limitation the 7 | rights to use, copy, modify, merge, publish, distribute, sublicense, and/or 8 | sell copies of the Software, and to permit persons to whom the Software is 9 | furnished to do so, subject to the following conditions: 10 | 11 | The above copyright notice and this permission notice shall be included in 12 | all copies or substantial portions of the Software. 13 | 14 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 17 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 20 | IN THE SOFTWARE. 21 | */ 22 | 23 | #if canImport(UIKit) 24 | 25 | import UIKit 26 | 27 | extension UIImage { 28 | /** 29 | Converts the image into an array of RGBA bytes. 30 | */ 31 | @nonobjc public func toByteArrayRGBA() -> [UInt8]? { 32 | return cgImage?.toByteArrayRGBA() 33 | } 34 | 35 | /** 36 | Creates a new UIImage from an array of RGBA bytes. 37 | */ 38 | @nonobjc public class func fromByteArrayRGBA(_ bytes: [UInt8], 39 | width: Int, 40 | height: Int, 41 | scale: CGFloat = 0, 42 | orientation: UIImage.Orientation = .up) -> UIImage? { 43 | if let cgImage = CGImage.fromByteArrayRGBA(bytes, width: width, height: height) { 44 | return UIImage(cgImage: cgImage, scale: scale, orientation: orientation) 45 | } else { 46 | return nil 47 | } 48 | } 49 | 50 | /** 51 | Creates a new UIImage from an array of grayscale bytes. 52 | */ 53 | @nonobjc public class func fromByteArrayGray(_ bytes: [UInt8], 54 | width: Int, 55 | height: Int, 56 | scale: CGFloat = 0, 57 | orientation: UIImage.Orientation = .up) -> UIImage? { 58 | if let cgImage = CGImage.fromByteArrayGray(bytes, width: width, height: height) { 59 | return UIImage(cgImage: cgImage, scale: scale, orientation: orientation) 60 | } else { 61 | return nil 62 | } 63 | } 64 | } 65 | 66 | #endif 67 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/DeepLabV3.mlmodel: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tbchen/BackgroundRemovalWithCoreMLSample/64713f5dd67ba43499bc968da899d336b189cb1e/BackgroundRemovalWithCoreMLSample/DeepLabV3.mlmodel -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/Info.plist: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | CFBundleDevelopmentRegion 6 | $(DEVELOPMENT_LANGUAGE) 7 | CFBundleExecutable 8 | $(EXECUTABLE_NAME) 9 | CFBundleIdentifier 10 | $(PRODUCT_BUNDLE_IDENTIFIER) 11 | CFBundleInfoDictionaryVersion 12 | 6.0 13 | CFBundleName 14 | $(PRODUCT_NAME) 15 | CFBundlePackageType 16 | APPL 17 | CFBundleShortVersionString 18 | 1.0 19 | CFBundleVersion 20 | 1 21 | LSRequiresIPhoneOS 22 | 23 | UILaunchStoryboardName 24 | LaunchScreen 25 | UIMainStoryboardFile 26 | Main 27 | UIRequiredDeviceCapabilities 28 | 29 | armv7 30 | 31 | UISupportedInterfaceOrientations 32 | 33 | UIInterfaceOrientationPortrait 34 | UIInterfaceOrientationLandscapeLeft 35 | UIInterfaceOrientationLandscapeRight 36 | 37 | UISupportedInterfaceOrientations~ipad 38 | 39 | UIInterfaceOrientationPortrait 40 | UIInterfaceOrientationPortraitUpsideDown 41 | UIInterfaceOrientationLandscapeLeft 42 | UIInterfaceOrientationLandscapeRight 43 | 44 | 45 | 46 | -------------------------------------------------------------------------------- /BackgroundRemovalWithCoreMLSample/ViewController.swift: -------------------------------------------------------------------------------- 1 | /* 2 | MIT License 3 | 4 | Copyright (c) 2019 tbchen 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. 23 | */ 24 | 25 | import UIKit 26 | 27 | class ViewController: UIViewController { 28 | 29 | // DeeplabV3 model from https://developer.apple.com/machine-learning/models/ 30 | let model = DeepLabV3() 31 | let context = CIContext(options: nil) 32 | 33 | @IBOutlet weak var inputImageView: UIImageView! 34 | @IBOutlet weak var outputImageView: UIImageView! 35 | 36 | override func viewDidLoad() { 37 | super.viewDidLoad() 38 | 39 | inputImageView.image = #imageLiteral(resourceName: "cat") 40 | outputImageView.image = removeBackground(image: inputImageView.image!) 41 | } 42 | 43 | private func removeBackground(image:UIImage) -> UIImage?{ 44 | let resizedImage = image.resized(to: CGSize(width: 513, height: 513)) 45 | if let pixelBuffer = resizedImage.pixelBuffer(width: Int(resizedImage.size.width), height: Int(resizedImage.size.height)){ 46 | if let outputImage = (try? model.prediction(image: pixelBuffer))?.semanticPredictions.image(min: 0, max: 1, axes: (0,0,1)), let outputCIImage = CIImage(image:outputImage){ 47 | if let maskImage = removeWhitePixels(image:outputCIImage), let resizedCIImage = CIImage(image: resizedImage), let compositedImage = composite(image: resizedCIImage, mask: maskImage){ 48 | return UIImage(ciImage: compositedImage).resized(to: CGSize(width: image.size.width, height: image.size.height)) 49 | } 50 | } 51 | } 52 | return nil 53 | } 54 | 55 | private func removeWhitePixels(image:CIImage) -> CIImage?{ 56 | let chromaCIFilter = chromaKeyFilter() 57 | chromaCIFilter?.setValue(image, forKey: kCIInputImageKey) 58 | return chromaCIFilter?.outputImage 59 | } 60 | 61 | private func composite(image:CIImage,mask:CIImage) -> CIImage?{ 62 | return CIFilter(name:"CISourceOutCompositing",parameters: 63 | [kCIInputImageKey: image,kCIInputBackgroundImageKey: mask])?.outputImage 64 | } 65 | 66 | // modified from https://developer.apple.com/documentation/coreimage/applying_a_chroma_key_effect 67 | private func chromaKeyFilter() -> CIFilter? { 68 | let size = 64 69 | var cubeRGB = [Float]() 70 | 71 | for z in 0 ..< size { 72 | let blue = CGFloat(z) / CGFloat(size-1) 73 | for y in 0 ..< size { 74 | let green = CGFloat(y) / CGFloat(size-1) 75 | for x in 0 ..< size { 76 | let red = CGFloat(x) / CGFloat(size-1) 77 | let brightness = getBrightness(red: red, green: green, blue: blue) 78 | let alpha: CGFloat = brightness == 1 ? 0 : 1 79 | cubeRGB.append(Float(red * alpha)) 80 | cubeRGB.append(Float(green * alpha)) 81 | cubeRGB.append(Float(blue * alpha)) 82 | cubeRGB.append(Float(alpha)) 83 | } 84 | } 85 | } 86 | 87 | let data = Data(buffer: UnsafeBufferPointer(start: &cubeRGB, count: cubeRGB.count)) 88 | 89 | let colorCubeFilter = CIFilter(name: "CIColorCube", parameters: ["inputCubeDimension": size, "inputCubeData": data]) 90 | return colorCubeFilter 91 | } 92 | 93 | // modified from https://developer.apple.com/documentation/coreimage/applying_a_chroma_key_effect 94 | private func getBrightness(red: CGFloat, green: CGFloat, blue: CGFloat) -> CGFloat { 95 | let color = UIColor(red: red, green: green, blue: blue, alpha: 1) 96 | var brightness: CGFloat = 0 97 | color.getHue(nil, saturation: nil, brightness: &brightness, alpha: nil) 98 | return brightness 99 | } 100 | } 101 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 tbchen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # BackgroundRemovalWithCoreMLSample 2 | 3 | This is an image background removal with Core ML and Core Image sample in Swift. 4 | 5 | ![Screen shot](ScreenShot.png) 6 | -------------------------------------------------------------------------------- /ScreenShot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tbchen/BackgroundRemovalWithCoreMLSample/64713f5dd67ba43499bc968da899d336b189cb1e/ScreenShot.png --------------------------------------------------------------------------------