From 1123097c92e87ab9cc5a90c8bab4e4d8f1c7d6ce Mon Sep 17 00:00:00 2001 From: Philipp Grosswiler Date: Sun, 3 Sep 2023 11:14:13 +0200 Subject: [PATCH] feat: add @capacitor-mlkit/selfie-segmentation package (#70) --- .../ios/Plugin.xcodeproj/project.pbxproj | 16 +++++++ .../Plugin/Classes/ProcessImageOptions.swift | 23 +++++++++ .../Plugin/Classes/ProcessImageResult.swift | 47 ++++++++++++++++++ .../ios/Plugin/SelfieSegmentation.swift | 48 +++++++++++++++++-- .../ios/Plugin/SelfieSegmentationPlugin.m | 2 +- .../ios/Plugin/SelfieSegmentationPlugin.swift | 39 ++++++++++++--- 6 files changed, 165 insertions(+), 10 deletions(-) create mode 100644 packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift create mode 100644 packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift diff --git a/packages/selfie-segmentation/ios/Plugin.xcodeproj/project.pbxproj b/packages/selfie-segmentation/ios/Plugin.xcodeproj/project.pbxproj index a6d581f..865a496 100644 --- a/packages/selfie-segmentation/ios/Plugin.xcodeproj/project.pbxproj +++ b/packages/selfie-segmentation/ios/Plugin.xcodeproj/project.pbxproj @@ -10,6 +10,8 @@ 03FC29A292ACC40490383A1F /* Pods_Plugin.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 3B2A61DA5A1F2DD4F959604D /* Pods_Plugin.framework */; }; 20C0B05DCFC8E3958A738AF2 /* Pods_PluginTests.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = F6753A823D3815DB436415E3 /* Pods_PluginTests.framework */; }; 2F98D68224C9AAE500613A4C /* SelfieSegmentation.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2F98D68124C9AAE400613A4C /* SelfieSegmentation.swift */; }; + 4A4E0DE52AA478A800BED263 /* ProcessImageOptions.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A4E0DE32AA478A800BED263 /* ProcessImageOptions.swift */; }; + 4A4E0DE62AA478A800BED263 /* ProcessImageResult.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A4E0DE42AA478A800BED263 /* ProcessImageResult.swift */; }; 50ADFF92201F53D600D50D53 /* Plugin.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 50ADFF88201F53D600D50D53 /* Plugin.framework */; }; 50ADFF97201F53D600D50D53 /* SelfieSegmentationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 50ADFF96201F53D600D50D53 /* SelfieSegmentationTests.swift */; }; 50ADFF99201F53D600D50D53 /* SelfieSegmentationPlugin.h in Headers */ = {isa = PBXBuildFile; fileRef = 50ADFF8B201F53D600D50D53 /* SelfieSegmentationPlugin.h */; settings = {ATTRIBUTES = (Public, ); }; }; @@ -31,6 +33,8 @@ /* Begin PBXFileReference section */ 2F98D68124C9AAE400613A4C /* SelfieSegmentation.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = SelfieSegmentation.swift; sourceTree = ""; }; 3B2A61DA5A1F2DD4F959604D /* Pods_Plugin.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_Plugin.framework; sourceTree = BUILT_PRODUCTS_DIR; }; + 4A4E0DE32AA478A800BED263 /* ProcessImageOptions.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = ProcessImageOptions.swift; path = Classes/ProcessImageOptions.swift; sourceTree = ""; }; + 4A4E0DE42AA478A800BED263 /* ProcessImageResult.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = ProcessImageResult.swift; path = Classes/ProcessImageResult.swift; sourceTree = ""; }; 50ADFF88201F53D600D50D53 /* Plugin.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Plugin.framework; sourceTree = BUILT_PRODUCTS_DIR; }; 50ADFF8B201F53D600D50D53 /* SelfieSegmentationPlugin.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = SelfieSegmentationPlugin.h; sourceTree = ""; }; 50ADFF8C201F53D600D50D53 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = ""; }; @@ -69,6 +73,15 @@ /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ + 4A4E0DE22AA4788400BED263 /* Classes */ = { + isa = PBXGroup; + children = ( + 4A4E0DE32AA478A800BED263 /* ProcessImageOptions.swift */, + 4A4E0DE42AA478A800BED263 /* ProcessImageResult.swift */, + ); + name = Classes; + sourceTree = ""; + }; 50ADFF7E201F53D600D50D53 = { isa = PBXGroup; children = ( @@ -92,6 +105,7 @@ 50ADFF8A201F53D600D50D53 /* Plugin */ = { isa = PBXGroup; children = ( + 4A4E0DE22AA4788400BED263 /* Classes */, 50E1A94720377CB70090CE1A /* SelfieSegmentationPlugin.swift */, 2F98D68124C9AAE400613A4C /* SelfieSegmentation.swift */, 50ADFF8B201F53D600D50D53 /* SelfieSegmentationPlugin.h */, @@ -364,8 +378,10 @@ isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( + 4A4E0DE62AA478A800BED263 /* ProcessImageResult.swift in Sources */, 50E1A94820377CB70090CE1A /* SelfieSegmentationPlugin.swift in Sources */, 2F98D68224C9AAE500613A4C /* SelfieSegmentation.swift in Sources */, + 4A4E0DE52AA478A800BED263 /* ProcessImageOptions.swift in Sources */, 50ADFFA82020EE4F00D50D53 /* SelfieSegmentationPlugin.m in Sources */, ); runOnlyForDeploymentPostprocessing = 0; diff --git a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift new file mode 100644 index 0000000..7dc31e7 --- /dev/null +++ b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift @@ -0,0 +1,23 @@ +import Foundation +import MLKitVision + +@objc class ProcessImageOptions: NSObject { + private var visionImage: VisionImage + private var enableRawSizeMask: Bool + + init( + visionImage: VisionImage, + enableRawSizeMask: Bool + ) { + self.visionImage = visionImage + self.enableRawSizeMask = enableRawSizeMask + } + + func getVisionImage() -> VisionImage { + return visionImage + } + + func shouldEnableRawSizeMask() -> Bool { + return enableRawSizeMask + } +} diff --git a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift new file mode 100644 index 0000000..ef8781c --- /dev/null +++ b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift @@ -0,0 +1,47 @@ +import Foundation +import Capacitor +import MLKitVision +import MLKitSegmentationSelfie + +@objc class ProcessImageResult: NSObject { + let segmentationMask: SegmentationMask + + init(segmentationMask: SegmentationMask) { + self.segmentationMask = segmentationMask + } + + func toJSObject() -> JSObject { + let (maskResult, maskWidth, maskHeight) = createMaskResult(mask: segmentationMask) + + var result = JSObject() + result["mask"] = maskResult + result["width"] = maskWidth + result["height"] = maskHeight + + return result + } + + private func createMaskResult(mask: SegmentationMask) -> (JSArray, Int, Int) { + var result = JSArray() + + let maskWidth = CVPixelBufferGetWidth(mask.buffer) + let maskHeight = CVPixelBufferGetHeight(mask.buffer) + + CVPixelBufferLockBaseAddress(mask.buffer, CVPixelBufferLockFlags.readOnly) + let maskBytesPerRow = CVPixelBufferGetBytesPerRow(mask.buffer) + var maskAddress = + CVPixelBufferGetBaseAddress(mask.buffer)!.bindMemory( + to: Float32.self, capacity: maskBytesPerRow * maskHeight) + + for _ in 0...(maskHeight - 1) { + for col in 0...(maskWidth - 1) { + // Gets the confidence of the pixel in the mask being in the foreground. + let foregroundConfidence: Float32 = maskAddress[col] + result.append(foregroundConfidence) + } + maskAddress += maskBytesPerRow / MemoryLayout.size + } + + return (result, maskWidth, maskHeight) + } +} diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift index 977f2f1..fc3e495 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift @@ -1,8 +1,50 @@ import Foundation +import MLKitVision +import MLKitSegmentationSelfie @objc public class SelfieSegmentation: NSObject { - @objc public func echo(_ value: String) -> String { - print(value) - return value + public let plugin: SelfieSegmentationPlugin + + init(plugin: SelfieSegmentationPlugin) { + self.plugin = plugin + } + + @objc func createVisionImageFromFilePath(_ path: String) -> VisionImage? { + guard let url = URL.init(string: path) else { + return nil + } + if FileManager.default.fileExists(atPath: url.path) { + guard let image = UIImage.init(contentsOfFile: url.path) else { + return nil + } + return VisionImage.init( + image: image + ) + } else { + return nil + } + } + + @objc func processImage(_ options: ProcessImageOptions, completion: @escaping (ProcessImageResult?, Error?) -> Void) { + let visionImage = options.getVisionImage() + let enableRawSizeMask = options.shouldEnableRawSizeMask() + + let selfieSegmenterOptions: SelfieSegmenterOptions = SelfieSegmenterOptions() + selfieSegmenterOptions.segmenterMode = .singleImage + selfieSegmenterOptions.shouldEnableRawSizeMask = enableRawSizeMask + + let segmenter = Segmenter.segmenter( + options: selfieSegmenterOptions + ) + + do { + let mask: SegmentationMask = try segmenter.results( + in: visionImage + ) + let result = ProcessImageResult(segmentationMask: mask) + completion(result, nil) + } catch let error { + completion(nil, error) + } } } diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.m b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.m index e850f82..6ec0c6b 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.m +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.m @@ -4,5 +4,5 @@ // Define the plugin using the CAP_PLUGIN Macro, and // each method the plugin supports using the CAP_PLUGIN_METHOD macro. CAP_PLUGIN(SelfieSegmentationPlugin, "SelfieSegmentation", - CAP_PLUGIN_METHOD(echo, CAPPluginReturnPromise); + CAP_PLUGIN_METHOD(processImage, CAPPluginReturnPromise); ) diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift index 69f9a98..1c35a18 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift @@ -7,12 +7,39 @@ import Capacitor */ @objc(SelfieSegmentationPlugin) public class SelfieSegmentationPlugin: CAPPlugin { - private let implementation = SelfieSegmentation() + public let tag = "SelfieSegmentation" + public let errorPathMissing = "path must be provided." + public let errorLoadImageFailed = "image could not be loaded." - @objc func echo(_ call: CAPPluginCall) { - let value = call.getString("value") ?? "" - call.resolve([ - "value": implementation.echo(value) - ]) + private var implementation: SelfieSegmentation? + + override public func load() { + implementation = SelfieSegmentation(plugin: self) + } + + @objc func processImage(_ call: CAPPluginCall) { + guard let path = call.getString("path") else { + call.reject(errorPathMissing) + return + } + let enableRawSizeMask = call.getBool("enableRawSizeMask", false) + + guard let visionImage = implementation?.createVisionImageFromFilePath(path) else { + call.reject(errorLoadImageFailed) + return + } + + let options = ProcessImageOptions(visionImage: visionImage, enableRawSizeMask: enableRawSizeMask) + + implementation?.processImage(options, completion: { result, error in + if let error = error { + CAPLog.print("[", self.tag, "] ", error) + call.reject(error.localizedDescription, nil, error) + return + } + if let result = result?.toJSObject() as? JSObject { + call.resolve(result) + } + }) } }