-
-
Notifications
You must be signed in to change notification settings - Fork 49
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add @capacitor-mlkit/selfie-segmentation package (#70)
- Loading branch information
Showing
6 changed files
with
165 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
23 changes: 23 additions & 0 deletions
23
packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } | ||
| } |
47 changes: 47 additions & 0 deletions
47
packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<Float32>.size | ||
| } | ||
|
|
||
| return (result, maskWidth, maskHeight) | ||
| } | ||
| } |
48 changes: 45 additions & 3 deletions
48
packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters