diff --git a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift index fcf830e..41b55e9 100644 --- a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift +++ b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift @@ -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 } } diff --git a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift index ef8781c..7a8e382 100644 --- a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift +++ b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageResult.swift @@ -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.size + + result["width"] = Int(image.size.width) + result["height"] = Int(image.size.height) } - return (result, maskWidth, maskHeight) + return result } } diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift index a8d4f28..b536b62 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift @@ -9,41 +9,191 @@ import MLKitSegmentationSelfie self.plugin = plugin } - @objc func createVisionImageFromFilePath(_ path: String) -> VisionImage? { + @objc func createImageFromFilePath(_ path: String) -> UIImage? { 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 - ) + return UIImage.init(contentsOfFile: url.path) } else { return nil } } + enum ProcessError: Error { + case createImageBuffer + } + + private var segmenter: Segmenter? + @objc func processImage(_ options: ProcessImageOptions, completion: @escaping (ProcessImageResult?, Error?) -> Void) { - let visionImage = options.getVisionImage() + let image = options.getImage() + let threshold = options.getConfidence() + + let visionImage = VisionImage.init(image: image) + visionImage.orientation = image.imageOrientation let selfieSegmenterOptions: SelfieSegmenterOptions = SelfieSegmenterOptions() selfieSegmenterOptions.segmenterMode = .singleImage - selfieSegmenterOptions.shouldEnableRawSizeMask = true - let segmenter = Segmenter.segmenter( + 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) + segmenter?.process(visionImage) { mask, error in + self.segmenter = nil + + guard error == nil, let mask = mask else { + return completion(nil, error) + } + + do { + guard let imageBuffer = self.createImageBuffer(from: image) else { + throw ProcessError.createImageBuffer + } + + self.applySegmentationMask( + mask: mask, to: imageBuffer, threshold: threshold + ) + + let image = self.createImage(from: imageBuffer) + let result = ProcessImageResult(image: image) + + completion(result, nil) + } catch { + completion(nil, error) + } + } + } + + func createImageBuffer(from image: UIImage) -> CVImageBuffer? { + guard let cgImage = image.cgImage else { return nil } + let width = cgImage.width + let height = cgImage.height + + var buffer: CVPixelBuffer? + CVPixelBufferCreate( + kCFAllocatorDefault, + width, + height, + kCVPixelFormatType_32BGRA, + nil, + &buffer) + guard let imageBuffer = buffer else { return nil } + + let flags = CVPixelBufferLockFlags(rawValue: 0) + CVPixelBufferLockBaseAddress(imageBuffer, flags) + let baseAddress = CVPixelBufferGetBaseAddress(imageBuffer) + let colorSpace = CGColorSpaceCreateDeviceRGB() + let bytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer) + let context = CGContext( + data: baseAddress, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: (CGImageAlphaInfo.premultipliedFirst.rawValue + | CGBitmapInfo.byteOrder32Little.rawValue)) + + if let context = context { + let rect = CGRect.init(x: 0, y: 0, width: width, height: height) + context.draw(cgImage, in: rect) + CVPixelBufferUnlockBaseAddress(imageBuffer, flags) + return imageBuffer + } else { + CVPixelBufferUnlockBaseAddress(imageBuffer, flags) + return nil + } + } + + // func createSampleBuffer(with imageBuffer: CVImageBuffer) -> CMSampleBuffer? { + // var timingInfo = CMSampleTimingInfo() + //// guard CMSampleBufferGetSampleTimingInfo(sampleBuffer, at: 0, timingInfoOut: &timingInfo) == 0 else { + //// return nil + //// } + // var outputSampleBuffer: CMSampleBuffer? + // var newFormatDescription: CMFormatDescription? + // CMVideoFormatDescriptionCreateForImageBuffer(allocator: nil, imageBuffer: imageBuffer, formatDescriptionOut: &newFormatDescription) + // guard let formatDescription = newFormatDescription else { + // return nil + // } + // CMSampleBufferCreateReadyWithImageBuffer(allocator: nil, imageBuffer: imageBuffer, formatDescription: formatDescription, sampleTiming: &timingInfo, sampleBufferOut: &outputSampleBuffer) + // guard let buffer = outputSampleBuffer else { + // return nil + // } + // return buffer + // } + + func createImage( + from imageBuffer: CVImageBuffer + ) -> UIImage { + let ciImage = CIImage(cvPixelBuffer: imageBuffer) + let context = CIContext(options: nil) + let cgImage = context.createCGImage(ciImage, from: ciImage.extent)! + return UIImage(cgImage: cgImage) + } + + func applySegmentationMask( + mask: SegmentationMask, to imageBuffer: CVImageBuffer, threshold: CGFloat + ) { + let bgraBytesPerPixel = 4 + + assert( + CVPixelBufferGetPixelFormatType(imageBuffer) == kCVPixelFormatType_32BGRA, + "Image buffer must have 32BGRA pixel format type") + + let width = CVPixelBufferGetWidth(mask.buffer) + let height = CVPixelBufferGetHeight(mask.buffer) + assert(CVPixelBufferGetWidth(imageBuffer) == width, "Width must match") + assert(CVPixelBufferGetHeight(imageBuffer) == height, "Height must match") + + let writeFlags = CVPixelBufferLockFlags(rawValue: 0) + CVPixelBufferLockBaseAddress(imageBuffer, writeFlags) + CVPixelBufferLockBaseAddress(mask.buffer, CVPixelBufferLockFlags.readOnly) + + let maskBytesPerRow = CVPixelBufferGetBytesPerRow(mask.buffer) + var maskAddress = + CVPixelBufferGetBaseAddress(mask.buffer)!.bindMemory( + to: Float32.self, capacity: maskBytesPerRow * height) + + let imageBytesPerRow = CVPixelBufferGetBytesPerRow(imageBuffer) + var imageAddress = CVPixelBufferGetBaseAddress(imageBuffer)!.bindMemory( + to: UInt8.self, capacity: imageBytesPerRow * height) + + for _ in 0...(height - 1) { + for col in 0...(width - 1) { + let pixelOffset = col * bgraBytesPerPixel + let blueOffset = pixelOffset + let greenOffset = pixelOffset + 1 + let redOffset = pixelOffset + 2 + let alphaOffset = pixelOffset + 3 + + let confidence: CGFloat = CGFloat(maskAddress[col]) + + if confidence >= threshold { + let red = CGFloat(imageAddress[redOffset]) + let green = CGFloat(imageAddress[greenOffset]) + let blue = CGFloat(imageAddress[blueOffset]) + // let alpha = CGFloat(imageAddress[alphaOffset]) + + imageAddress[redOffset] = UInt8(red * confidence) + imageAddress[greenOffset] = UInt8(green * confidence) + imageAddress[blueOffset] = UInt8(blue * confidence) + imageAddress[alphaOffset] = UInt8(0xff) + } else { + imageAddress[redOffset] = UInt8(0x00) + imageAddress[greenOffset] = UInt8(0x00) + imageAddress[blueOffset] = UInt8(0x00) + imageAddress[alphaOffset] = UInt8(0x00) + } + } + + imageAddress += imageBytesPerRow / MemoryLayout.size + maskAddress += maskBytesPerRow / MemoryLayout.size } + + CVPixelBufferUnlockBaseAddress(imageBuffer, writeFlags) + CVPixelBufferUnlockBaseAddress(mask.buffer, CVPixelBufferLockFlags.readOnly) } } diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift index e0b20ea..8175d3b 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift @@ -8,9 +8,12 @@ import Capacitor @objc(SelfieSegmentationPlugin) public class SelfieSegmentationPlugin: CAPPlugin { public let tag = "SelfieSegmentation" + public let errorPathMissing = "path must be provided." public let errorLoadImageFailed = "image could not be loaded." + public let defaultConfidence: Float = 0.9 + private var implementation: SelfieSegmentation? override public func load() { @@ -23,12 +26,20 @@ public class SelfieSegmentationPlugin: CAPPlugin { return } - guard let visionImage = implementation?.createVisionImageFromFilePath(path) else { + let width = call.getInt("width") + let height = call.getInt("height") + + let confidence = call.getFloat("confidence", defaultConfidence) + + guard let image = implementation?.createImageFromFilePath(path) else { call.reject(errorLoadImageFailed) return } - let options = ProcessImageOptions(visionImage: visionImage) + let options = ProcessImageOptions(image: image, + width: width, + height: height, + confidence: CGFloat(confidence)) implementation?.processImage(options, completion: { result, error in if let error = error { @@ -36,8 +47,9 @@ public class SelfieSegmentationPlugin: CAPPlugin { call.reject(error.localizedDescription, nil, error) return } - if let result = result?.toJSObject() as? JSObject { - call.resolve(result) + + if let result = result { + call.resolve(result.toJSObject()) } }) }