From bbadc5b5ba7853071f25ea82ed83262ecf09d5ef Mon Sep 17 00:00:00 2001 From: Philipp Grosswiler Date: Sat, 9 Sep 2023 11:52:55 +0200 Subject: [PATCH] feat: add @capacitor-mlkit/selfie-segmentation package (#73) --- packages/selfie-segmentation/README.md | 20 ++--- .../SelfieSegmentation.java | 81 ++++++++++++++++--- .../SelfieSegmentationPlugin.java | 11 ++- .../classes/ProcessImageOptions.java | 38 +++++++-- .../classes/ProcessImageResult.java | 38 +++------ .../Plugin/Classes/ProcessImageOptions.swift | 9 +-- .../ios/Plugin/SelfieSegmentation.swift | 3 +- .../ios/Plugin/SelfieSegmentationPlugin.swift | 3 +- .../selfie-segmentation/src/definitions.ts | 32 +++++--- 9 files changed, 160 insertions(+), 75 deletions(-) diff --git a/packages/selfie-segmentation/README.md b/packages/selfie-segmentation/README.md index 8b3b719..d253246 100644 --- a/packages/selfie-segmentation/README.md +++ b/packages/selfie-segmentation/README.md @@ -71,19 +71,21 @@ Only available on Android and iOS. #### ProcessImageResult -| Prop | Type | Description | Since | -| ------------ | --------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----- | -| **`mask`** | number[] | Returns a mask that indicates the foreground and background segmentation. This mask’s dimensions could vary, depending on whether a raw size mask is requested via options. | 5.2.0 | -| **`width`** | number | Returns the width of the mask. | 5.2.0 | -| **`height`** | number | Returns the height of the mask. | 5.2.0 | +| Prop | Type | Description | Since | +| ------------ | ------------------- | ------------------------------------- | ----- | +| **`path`** | string | The path to the segmented image file. | 5.2.0 | +| **`width`** | number | Returns the width of the image file. | 5.2.0 | +| **`height`** | number | Returns the height of the image file. | 5.2.0 | #### ProcessImageOptions -| Prop | Type | Description | Default | Since | -| ----------------------- | -------------------- | ----------------------------------------------------------------------------------- | ------------------ | ----- | -| **`path`** | string | The local path to the image file. | | 5.2.0 | -| **`enableRawSizeMask`** | boolean | Asks the segmenter to return the raw size mask which matches the model output size. | false | 5.2.0 | +| Prop | Type | Description | Default | Since | +| ---------------- | ------------------- | ----------------------------------------------------------------------------------------- | ---------------- | ----- | +| **`path`** | string | The local path to the image file. | | 5.2.0 | +| **`width`** | number | Scale the image to this width. If no `height` is given, it will respect the aspect ratio. | | 5.2.0 | +| **`height`** | number | Scale the image to this height. If no `width` is given, it will respect the aspect ratio. | | 5.2.0 | +| **`confidence`** | number | Sets the confidence threshold. | 0.9 | 5.2.0 | diff --git a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentation.java b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentation.java index c28dddd..3c7e6e5 100644 --- a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentation.java +++ b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentation.java @@ -1,5 +1,7 @@ package io.capawesome.capacitorjs.plugins.mlkit.selfiesegmentation; +import android.annotation.SuppressLint; +import android.graphics.Bitmap; import android.net.Uri; import androidx.annotation.NonNull; import androidx.annotation.Nullable; @@ -9,6 +11,14 @@ import com.google.mlkit.vision.segmentation.selfie.SelfieSegmenterOptions; import io.capawesome.capacitorjs.plugins.mlkit.selfiesegmentation.classes.ProcessImageOptions; import io.capawesome.capacitorjs.plugins.mlkit.selfiesegmentation.classes.ProcessImageResult; +import java.io.File; +import java.io.FileOutputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.text.SimpleDateFormat; +import java.util.Date; +import java.util.Objects; public class SelfieSegmentation { @@ -30,27 +40,79 @@ public InputImage createInputImageFromFilePath(@NonNull String path) { public void processImage(ProcessImageOptions options, ProcessImageResultCallback callback) { InputImage inputImage = options.getInputImage(); - boolean enableRawSizeMask = options.isRawSizeMaskEnabled(); + Float threshold = options.getConfidence(); SelfieSegmenterOptions.Builder builder = new SelfieSegmenterOptions.Builder(); builder.setDetectorMode(SelfieSegmenterOptions.SINGLE_IMAGE_MODE); - if (enableRawSizeMask) { - builder.enableRawSizeMask(); - } SelfieSegmenterOptions selfieSegmenterOptions = builder.build(); final Segmenter segmenter = Segmentation.getClient(selfieSegmenterOptions); + plugin .getActivity() .runOnUiThread( - () -> { + () -> segmenter .process(inputImage) .addOnSuccessListener( - mask -> { + segmentationMask -> { segmenter.close(); - ProcessImageResult result = new ProcessImageResult(mask); - callback.success(result); + + ByteBuffer mask = segmentationMask.getBuffer(); + + Bitmap bitmap = inputImage.getBitmapInternal(); + Objects.requireNonNull(bitmap).setHasAlpha(true); + + ByteBuffer pixels = ByteBuffer.allocateDirect(bitmap.getAllocationByteCount()); + bitmap.copyPixelsToBuffer(pixels); + + final boolean bigEndian = pixels.order() == ByteOrder.BIG_ENDIAN; + final int ALPHA = bigEndian ? 3 : 0; + final int RED = bigEndian ? 2 : 1; + final int GREEN = bigEndian ? 1 : 2; + final int BLUE = bigEndian ? 0 : 3; + + for (int i = 0; i < pixels.capacity() >> 2; i++) { + float confidence = mask.getFloat(); + + if (confidence >= threshold) { + byte red = pixels.get((i << 2) + RED); + byte green = pixels.get((i << 2) + GREEN); + byte blue = pixels.get((i << 2) + BLUE); + + pixels.put((i << 2) + ALPHA, (byte) (0xff)); + pixels.put((i << 2) + RED, (byte) (red * confidence)); + pixels.put((i << 2) + GREEN, (byte) (green * confidence)); + pixels.put((i << 2) + BLUE, (byte) (blue * confidence)); + } else { + pixels.putInt(i << 2, 0x00000000); // transparent + } + } + + bitmap.copyPixelsFromBuffer(pixels.rewind()); + + // Create an image file name + @SuppressLint("SimpleDateFormat") + String timeStamp = new SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date()); + String imageFileName = "PNG_" + timeStamp + "_"; + + try { + File image = File.createTempFile(imageFileName, ".png"); + + OutputStream stream = new FileOutputStream(image); + bitmap.compress(Bitmap.CompressFormat.PNG, 100, stream); + stream.close(); + + ProcessImageResult result = new ProcessImageResult( + image.getAbsolutePath(), + bitmap.getWidth(), + bitmap.getHeight() + ); + + callback.success(result); + } catch (Exception exception) { + callback.error(exception); + } } ) .addOnCanceledListener( @@ -64,8 +126,7 @@ public void processImage(ProcessImageOptions options, ProcessImageResultCallback segmenter.close(); callback.error(exception); } - ); - } + ) ); } } diff --git a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentationPlugin.java b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentationPlugin.java index 628c13f..533c0fe 100644 --- a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentationPlugin.java +++ b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/SelfieSegmentationPlugin.java @@ -13,10 +13,13 @@ public class SelfieSegmentationPlugin extends Plugin { public static final String TAG = "SelfieSegmentation"; + public static final String ERROR_PROCESS_IMAGE_CANCELED = "processImage canceled."; public static final String ERROR_PATH_MISSING = "path must be provided."; public static final String ERROR_LOAD_IMAGE_FAILED = "image could not be loaded."; + public static final float CONFIDENCE = 0.9f; + private SelfieSegmentation implementation; @Override @@ -36,14 +39,18 @@ public void processImage(PluginCall call) { call.reject(ERROR_PATH_MISSING); return; } - Boolean enableRawSizeMask = call.getBoolean("enableRawSizeMask", false); + + Integer width = call.getInt("width", null); + Integer height = call.getInt("height", null); + + Float confidence = call.getFloat("confidence", CONFIDENCE); InputImage image = implementation.createInputImageFromFilePath(path); if (image == null) { call.reject(ERROR_LOAD_IMAGE_FAILED); return; } - ProcessImageOptions options = new ProcessImageOptions(image, enableRawSizeMask); + ProcessImageOptions options = new ProcessImageOptions(image, width, height, confidence); implementation.processImage( options, diff --git a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageOptions.java b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageOptions.java index 8c6ded0..79755ab 100644 --- a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageOptions.java +++ b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageOptions.java @@ -1,23 +1,47 @@ package io.capawesome.capacitorjs.plugins.mlkit.selfiesegmentation.classes; +import android.graphics.Bitmap; import com.google.mlkit.vision.common.InputImage; +import java.util.Objects; public class ProcessImageOptions { - private InputImage inputImage; + private final InputImage inputImage; - private boolean enableRawSizeMask; + private final Float confidence; - public ProcessImageOptions(InputImage inputImage, boolean enableRawSizeMask) { - this.inputImage = inputImage; - this.enableRawSizeMask = enableRawSizeMask; + public ProcessImageOptions(InputImage inputImage, Integer width, Integer height, Float confidence) { + this.inputImage = scaledImage(inputImage, width, height); + + this.confidence = confidence; } public InputImage getInputImage() { return inputImage; } - public boolean isRawSizeMaskEnabled() { - return enableRawSizeMask; + public Float getConfidence() { + return confidence; + } + + private InputImage scaledImage(InputImage inputImage, Integer width, Integer height) { + float scaleX = (width != null) ? width * 1f / inputImage.getWidth() : 0f; + float scaleY = (height != null) ? height * 1f / inputImage.getHeight() : 0f; + + if (scaleX > 0f || scaleY > 0f) { + if (scaleX > 0f && scaleY == 0f) scaleY = scaleX; else if (scaleY > 0f && scaleX == 0f) scaleX = scaleY; + + return InputImage.fromBitmap( + Bitmap.createScaledBitmap( + Objects.requireNonNull(inputImage.getBitmapInternal()), + (int) (inputImage.getWidth() * scaleX), + (int) (inputImage.getHeight() * scaleY), + false + ), + inputImage.getRotationDegrees() + ); + } + + return inputImage; } } diff --git a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageResult.java b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageResult.java index 2a2c374..ffe3363 100644 --- a/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageResult.java +++ b/packages/selfie-segmentation/android/src/main/java/io/capawesome/capacitorjs/plugins/mlkit/selfiesegmentation/classes/ProcessImageResult.java @@ -1,41 +1,29 @@ package io.capawesome.capacitorjs.plugins.mlkit.selfiesegmentation.classes; -import com.getcapacitor.JSArray; import com.getcapacitor.JSObject; -import com.google.mlkit.vision.segmentation.SegmentationMask; -import java.nio.ByteBuffer; -import org.json.JSONException; public class ProcessImageResult { - private SegmentationMask segmentationMask; + private final String imagePath; - public ProcessImageResult(SegmentationMask segmentationMask) { - this.segmentationMask = segmentationMask; - } + private final int width; + private final int height; - public JSObject toJSObject() throws JSONException { - JSArray maskResult = this.createMaskResult(); + public ProcessImageResult(String imagePath, int width, int height) { + this.imagePath = imagePath; - JSObject result = new JSObject(); - result.put("mask", maskResult); - result.put("width", segmentationMask.getWidth()); - result.put("height", segmentationMask.getHeight()); - return result; + this.width = width; + this.height = height; } - private JSArray createMaskResult() throws JSONException { - JSArray result = new JSArray(); + public JSObject toJSObject() { + JSObject result = new JSObject(); + + result.put("path", imagePath); - ByteBuffer mask = segmentationMask.getBuffer(); - int maskWidth = segmentationMask.getWidth(); - int maskHeight = segmentationMask.getHeight(); + result.put("width", width); + result.put("height", height); - for (int y = 0; y < maskHeight; y++) { - for (int x = 0; x < maskWidth; x++) { - result.put(mask.getFloat()); - } - } return result; } } diff --git a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift index 7dc31e7..fcf830e 100644 --- a/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift +++ b/packages/selfie-segmentation/ios/Plugin/Classes/ProcessImageOptions.swift @@ -3,21 +3,14 @@ import MLKitVision @objc class ProcessImageOptions: NSObject { private var visionImage: VisionImage - private var enableRawSizeMask: Bool init( - visionImage: VisionImage, - enableRawSizeMask: Bool + visionImage: VisionImage ) { self.visionImage = visionImage - self.enableRawSizeMask = enableRawSizeMask } func getVisionImage() -> VisionImage { return visionImage } - - func shouldEnableRawSizeMask() -> Bool { - return enableRawSizeMask - } } diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift index fc3e495..a8d4f28 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentation.swift @@ -27,11 +27,10 @@ import MLKitSegmentationSelfie @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 + selfieSegmenterOptions.shouldEnableRawSizeMask = true let segmenter = Segmenter.segmenter( options: selfieSegmenterOptions diff --git a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift index 1c35a18..e0b20ea 100644 --- a/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift +++ b/packages/selfie-segmentation/ios/Plugin/SelfieSegmentationPlugin.swift @@ -22,14 +22,13 @@ public class SelfieSegmentationPlugin: CAPPlugin { 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) + let options = ProcessImageOptions(visionImage: visionImage) implementation?.processImage(options, completion: { result, error in if let error = error { diff --git a/packages/selfie-segmentation/src/definitions.ts b/packages/selfie-segmentation/src/definitions.ts index 077441e..1e8d6b6 100644 --- a/packages/selfie-segmentation/src/definitions.ts +++ b/packages/selfie-segmentation/src/definitions.ts @@ -21,36 +21,48 @@ export interface ProcessImageOptions { path: string; /** - * Asks the segmenter to return the raw size mask which matches the model output size. + * Scale the image to this width. + * If no `height` is given, it will respect the aspect ratio. * * @since 5.2.0 - * @default false */ - enableRawSizeMask?: boolean; + width?: number; + /** + * Scale the image to this height. + * If no `width` is given, it will respect the aspect ratio. + * + * @since 5.2.0 + */ + height?: number; + + /** + * Sets the confidence threshold. + * + * @since 5.2.0 + * @default 0.9 + */ + confidence?: number; } /** * @since 5.2.0 - * @see https://developers.google.com/android/reference/com/google/mlkit/vision/segmentation/SegmentationMask */ export interface ProcessImageResult { /** - * Returns a mask that indicates the foreground and background segmentation. - * - * This mask’s dimensions could vary, depending on whether a raw size mask is requested via options. + * The path to the segmented image file. * * @since 5.2.0 */ - mask: number[]; + path: string; /** - * Returns the width of the mask. + * Returns the width of the image file. * * @since 5.2.0 */ width: number; /** - * Returns the height of the mask. + * Returns the height of the image file. * * @since 5.2.0 */