diff --git a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imagesegmenter/ImageSegmenter.java b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imagesegmenter/ImageSegmenter.java index 8d07b7c68..2ef1b57d8 100644 --- a/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imagesegmenter/ImageSegmenter.java +++ b/mediapipe/tasks/java/com/google/mediapipe/tasks/vision/imagesegmenter/ImageSegmenter.java @@ -47,10 +47,14 @@ import java.util.Optional; /** * Performs image segmentation on images. * - *

Note that, unlike other vision tasks, the output of ImageSegmenter is provided through a - * user-defined callback function even for the synchronous API. This makes it possible for - * ImageSegmenter to return the output masks without any copy. {@link ResultListener} must be set in - * the {@link ImageSegmenterOptions} for all {@link RunningMode}. + *

Note that, in addition to the standard segmentation API, {@link segment} and {@link + * segmentForVideo}, that take an input image and return the outputs, but involves deep copy of the + * returns, ImageSegmenter also supports the callback API, {@link segmentWithResultListener} and + * {@link segmentForVideoWithResultListener}, which allow you to access the outputs through zero + * copy. + * + *

The callback API is available for all {@link RunningMode} in ImageSegmenter. Set {@link + * ResultListener} in {@link ImageSegmenterOptions} properly to use the callback API. * *

The API expects a TFLite model with,TFLite Model Metadata.. @@ -85,6 +89,8 @@ public final class ImageSegmenter extends BaseVisionTaskApi { private static final String TASK_GRAPH_NAME = "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph"; + private boolean hasResultListener = false; + /** * Creates an {@link ImageSegmenter} instance from an {@link ImageSegmenterOptions}. * @@ -116,8 +122,19 @@ public final class ImageSegmenter extends BaseVisionTaskApi { int imageListSize = PacketGetter.getImageListSize(packets.get(GROUPED_SEGMENTATION_OUT_STREAM_INDEX)); ByteBuffer[] buffersArray = new ByteBuffer[imageListSize]; + // If resultListener is not provided, the resulted MPImage is deep copied from mediapipe + // graph. If provided, the result MPImage is wrapping the mediapipe packet memory. + if (!segmenterOptions.resultListener().isPresent()) { + for (int i = 0; i < imageListSize; i++) { + buffersArray[i] = + ByteBuffer.allocateDirect( + width * height * (imageFormat == MPImage.IMAGE_FORMAT_VEC32F1 ? 4 : 1)); + } + } if (!PacketGetter.getImageList( - packets.get(GROUPED_SEGMENTATION_OUT_STREAM_INDEX), buffersArray, false)) { + packets.get(GROUPED_SEGMENTATION_OUT_STREAM_INDEX), + buffersArray, + !segmenterOptions.resultListener().isPresent())) { throw new MediaPipeException( MediaPipeException.StatusCode.INTERNAL.ordinal(), "There is an error getting segmented masks. It usually results from incorrect" @@ -143,7 +160,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi { .build(); } }); - handler.setResultListener(segmenterOptions.resultListener()); + segmenterOptions.resultListener().ifPresent(handler::setResultListener); segmenterOptions.errorListener().ifPresent(handler::setErrorListener); TaskRunner runner = TaskRunner.create( @@ -158,7 +175,8 @@ public final class ImageSegmenter extends BaseVisionTaskApi { .setEnableFlowLimiting(segmenterOptions.runningMode() == RunningMode.LIVE_STREAM) .build(), handler); - return new ImageSegmenter(runner, segmenterOptions.runningMode()); + return new ImageSegmenter( + runner, segmenterOptions.runningMode(), segmenterOptions.resultListener().isPresent()); } /** @@ -168,16 +186,17 @@ public final class ImageSegmenter extends BaseVisionTaskApi { * @param taskRunner a {@link TaskRunner}. * @param runningMode a mediapipe vision task {@link RunningMode}. */ - private ImageSegmenter(TaskRunner taskRunner, RunningMode runningMode) { + private ImageSegmenter( + TaskRunner taskRunner, RunningMode runningMode, boolean hasResultListener) { super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME); + this.hasResultListener = hasResultListener; } /** * Performs image segmentation on the provided single image with default image processing options, - * i.e. without any rotation applied, and the results will be available via the {@link - * ResultListener} provided in the {@link ImageSegmenterOptions}. Only use this method when the - * {@link ImageSegmenter} is created with {@link RunningMode.IMAGE}. TODO update java - * doc for input image format. + * i.e. without any rotation applied. Only use this method when the {@link ImageSegmenter} is + * created with {@link RunningMode.IMAGE}. TODO update java doc for input image + * format. * *

{@link ImageSegmenter} supports the following color space types: * @@ -186,19 +205,19 @@ public final class ImageSegmenter extends BaseVisionTaskApi { * * * @param image a MediaPipe {@link MPImage} object for processing. - * @throws MediaPipeException if there is an internal error. + * @throws MediaPipeException if there is an internal error. Or if {@link ImageSegmenter} is + * created with a {@link ResultListener}. */ - public void segment(MPImage image) { - segment(image, ImageProcessingOptions.builder().build()); + public ImageSegmenterResult segment(MPImage image) { + return segment(image, ImageProcessingOptions.builder().build()); } /** - * Performs image segmentation on the provided single image, and the results will be available via - * the {@link ResultListener} provided in the {@link ImageSegmenterOptions}. Only use this method - * when the {@link ImageSegmenter} is created with {@link RunningMode.IMAGE}. TODO - * update java doc for input image format. + * Performs image segmentation on the provided single image. Only use this method when the {@link + * ImageSegmenter} is created with {@link RunningMode.IMAGE}. TODO update java doc + * for input image format. * - *

{@link HandLandmarker} supports the following color space types: + *

{@link ImageSegmenter} supports the following color space types: * *