-
-
Notifications
You must be signed in to change notification settings - Fork 48
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
4 changed files
with
253 additions
and
59 deletions.
There are no files selected for viewing
59 changes: 53 additions & 6 deletions
59
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 |
---|---|---|
@@ -1,16 +1,63 @@ | ||
import Foundation | ||
import MLKitVision | ||
|
||
extension UIImage { | ||
public func scaledImage(width: Int?, height: Int?) -> UIImage { | ||
let newWidth: CGFloat | ||
let newHeight: CGFloat | ||
|
||
if let width = width { | ||
newWidth = CGFloat(width) | ||
if let height = height { | ||
newHeight = CGFloat(height) | ||
} else { | ||
let scaleFactor = newWidth / self.size.width | ||
newHeight = self.size.height * scaleFactor | ||
} | ||
} else | ||
if let height = height { | ||
newHeight = CGFloat(height) | ||
if let width = width { | ||
newWidth = CGFloat(width) | ||
} else { | ||
let scaleFactor = newHeight / self.size.height | ||
newWidth = self.size.width * scaleFactor | ||
} | ||
} else { | ||
return self | ||
} | ||
|
||
let newSize = CGSize(width: newWidth, height: newHeight) | ||
|
||
if newSize.width >= size.width && newSize.height >= size.height { | ||
return self | ||
} | ||
|
||
UIGraphicsBeginImageContextWithOptions(newSize, false, scale) | ||
defer { UIGraphicsEndImageContext() } | ||
draw(in: CGRect(origin: .zero, size: newSize)) | ||
return UIGraphicsGetImageFromCurrentImageContext() ?? self | ||
} | ||
} | ||
|
||
@objc class ProcessImageOptions: NSObject { | ||
private var visionImage: VisionImage | ||
private var image: UIImage | ||
private var confidence: CGFloat | ||
|
||
init( | ||
visionImage: VisionImage | ||
image: UIImage, | ||
width: Int?, | ||
height: Int?, | ||
confidence: CGFloat | ||
) { | ||
self.visionImage = visionImage | ||
self.image = image.scaledImage(width: width, height: height) | ||
self.confidence = confidence | ||
} | ||
|
||
func getImage() -> UIImage { | ||
return image | ||
} | ||
|
||
func getVisionImage() -> VisionImage { | ||
return visionImage | ||
func getConfidence() -> CGFloat { | ||
return confidence | ||
} | ||
} |
47 changes: 16 additions & 31 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 |
---|---|---|
@@ -1,47 +1,32 @@ | ||
import Foundation | ||
import Capacitor | ||
import MLKitVision | ||
import MLKitSegmentationSelfie | ||
|
||
@objc class ProcessImageResult: NSObject { | ||
let segmentationMask: SegmentationMask | ||
let image: UIImage | ||
|
||
init(segmentationMask: SegmentationMask) { | ||
self.segmentationMask = segmentationMask | ||
init(image: UIImage) { | ||
self.image = image | ||
} | ||
|
||
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() | ||
if let data = image.pngData() { | ||
do { | ||
let path = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask)[0] | ||
let name = "photo-"+UUID().uuidString+".png" | ||
let url = path.appendingPathComponent(name) | ||
try data.write(to: url) | ||
|
||
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) | ||
result["path"] = url.absoluteString | ||
} catch { | ||
result["path"] = "data:image/png;base64," + data.base64EncodedString() | ||
} | ||
maskAddress += maskBytesPerRow / MemoryLayout<Float32>.size | ||
|
||
result["width"] = Int(image.size.width) | ||
result["height"] = Int(image.size.height) | ||
} | ||
|
||
return (result, maskWidth, maskHeight) | ||
return result | ||
} | ||
} |
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