From 3fbb2b002bb6e4bc80fa62d7dc4d34d1df1fc287 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 21 Sep 2022 03:23:04 -0700 Subject: [PATCH 01/15] Added image segmenter implementation files --- mediapipe/python/BUILD | 2 + mediapipe/tasks/python/components/BUILD | 28 +++ .../python/components/segmenter_options.py | 78 +++++++ mediapipe/tasks/python/test/vision/BUILD | 38 +++- .../test/vision/image_segmenter_test.py | 118 ++++++++++ mediapipe/tasks/python/vision/BUILD | 19 ++ .../tasks/python/vision/image_segmenter.py | 205 ++++++++++++++++++ 7 files changed, 487 insertions(+), 1 deletion(-) create mode 100644 mediapipe/tasks/python/components/BUILD create mode 100644 mediapipe/tasks/python/components/segmenter_options.py create mode 100644 mediapipe/tasks/python/test/vision/image_segmenter_test.py create mode 100644 mediapipe/tasks/python/vision/image_segmenter.py diff --git a/mediapipe/python/BUILD b/mediapipe/python/BUILD index 3a4a90b4..331ee836 100644 --- a/mediapipe/python/BUILD +++ b/mediapipe/python/BUILD @@ -86,6 +86,8 @@ cc_library( name = "builtin_task_graphs", deps = [ "//mediapipe/tasks/cc/vision/object_detector:object_detector_graph", + "//mediapipe/tasks/cc/vision/image_classification:image_classifier_graph", + "//mediapipe/tasks/cc/vision/image_segmenter:image_segmenter_graph", ], ) diff --git a/mediapipe/tasks/python/components/BUILD b/mediapipe/tasks/python/components/BUILD new file mode 100644 index 00000000..eb8714a9 --- /dev/null +++ b/mediapipe/tasks/python/components/BUILD @@ -0,0 +1,28 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Placeholder for internal Python strict library compatibility macro. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +py_library( + name = "segmenter_options", + srcs = ["segmenter_options.py"], + deps = [ + "//mediapipe/tasks/cc/components:segmenter_options_py_pb2", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) diff --git a/mediapipe/tasks/python/components/segmenter_options.py b/mediapipe/tasks/python/components/segmenter_options.py new file mode 100644 index 00000000..5b94a256 --- /dev/null +++ b/mediapipe/tasks/python/components/segmenter_options.py @@ -0,0 +1,78 @@ +# Copyright 2022 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Segmenter options data class.""" + +import dataclasses +import enum +from typing import Any, Optional + +from mediapipe.tasks.cc.components import segmenter_options_pb2 +from mediapipe.tasks.python.core.optional_dependencies import doc_controls + +_SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions + + +class OutputType(enum.Enum): + UNSPECIFIED = 0 + CATEGORY_MASK = 1 + CONFIDENCE_MASK = 2 + + +class Activation(enum.Enum): + NONE = 0 + SIGMOID = 1 + SOFTMAX = 2 + + +@dataclasses.dataclass +class SegmenterOptions: + """Options for segmentation processor. + Attributes: + output_type: The output mask type allows specifying the type of + post-processing to perform on the raw model results. + activation: Activation function to apply to input tensor. + """ + + output_type: Optional[OutputType] = OutputType.CATEGORY_MASK + activation: Optional[Activation] = Activation.NONE + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _SegmenterOptionsProto: + """Generates a protobuf object to pass to the C++ layer.""" + return _SegmenterOptionsProto( + output_type=self.output_type.value, + activation=self.activation.value + ) + + @classmethod + @doc_controls.do_not_generate_docs + def create_from_pb2( + cls, pb2_obj: _SegmenterOptionsProto) -> "SegmenterOptions": + """Creates a `SegmenterOptions` object from the given protobuf object.""" + return SegmenterOptions( + output_type=OutputType(pb2_obj.output_type), + activation=Activation(pb2_obj.output_type) + ) + + def __eq__(self, other: Any) -> bool: + """Checks if this object is equal to the given object. + Args: + other: The object to be compared with. + Returns: + True if the objects are equal. + """ + if not isinstance(other, SegmenterOptions): + return False + + return self.to_pb2().__eq__(other.to_pb2()) diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index bb495338..403e00a3 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -18,4 +18,40 @@ package(default_visibility = ["//mediapipe/tasks:internal"]) licenses(["notice"]) -# TODO: This test fails in OSS +py_test( + name = "object_detector_test", + srcs = ["object_detector_test.py"], + data = [ + "//mediapipe/tasks/testdata/vision:test_images", + "//mediapipe/tasks/testdata/vision:test_models", + ], + deps = [ + # build rule placeholder: numpy dep, + "//mediapipe/tasks/python/components/containers:bounding_box", + "//mediapipe/tasks/python/components/containers:category", + "//mediapipe/tasks/python/components/containers:detections", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/test:test_util", + "//mediapipe/tasks/python/vision:object_detector", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + "@absl_py//absl/testing:parameterized", + ], +) + +py_test( + name = "image_segmenter_test", + srcs = ["image_segmenter_test.py"], + data = [ + "//mediapipe/tasks/testdata/vision:test_images", + "//mediapipe/tasks/testdata/vision:test_models", + ], + deps = [ + # build rule placeholder: numpy dep, + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/test:test_util", + "//mediapipe/tasks/python/components:segmenter_options", + "//mediapipe/tasks/python/vision:image_segmenter", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + "@absl_py//absl/testing:parameterized", + ], +) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py new file mode 100644 index 00000000..704b0fc5 --- /dev/null +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -0,0 +1,118 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for image segmenter.""" + +import enum + +from absl.testing import absltest +from absl.testing import parameterized +import numpy as np + +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.tasks.python.components import segmenter_options +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.test import test_util +from mediapipe.tasks.python.vision import image_segmenter +from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module + +_BaseOptions = base_options_module.BaseOptions +_Image = image_module.Image +_OutputType = segmenter_options.OutputType +_Activation = segmenter_options.Activation +_ImageSegmenter = image_segmenter.ImageSegmenter +_ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions +_RUNNING_MODE = running_mode_module.VisionTaskRunningMode + +_MODEL_FILE = 'deeplabv3.tflite' +_IMAGE_FILE = 'segmentation_input_rotation0.jpg' +_SEGMENTATION_FILE = 'segmentation_golden_rotation0.png' +_MASK_MAGNIFICATION_FACTOR = 10 +_MATCH_PIXELS_THRESHOLD = 0.01 + + +class ModelFileType(enum.Enum): + FILE_CONTENT = 1 + FILE_NAME = 2 + + +class ImageSegmenterTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.test_image = test_util.read_test_image( + test_util.get_test_data_path(_IMAGE_FILE)) + self.model_path = test_util.get_test_data_path(_MODEL_FILE) + + def test_create_from_file_succeeds_with_valid_model_path(self): + # Creates with default option and valid model file successfully. + with _ImageSegmenter.create_from_model_path(self.model_path) as segmenter: + self.assertIsInstance(segmenter, _ImageSegmenter) + + def test_create_from_options_succeeds_with_valid_model_path(self): + # Creates with options containing model file successfully. + base_options = _BaseOptions(file_name=self.model_path) + options = _ImageSegmenterOptions(base_options=base_options) + with _ImageSegmenter.create_from_options(options) as segmenter: + self.assertIsInstance(segmenter, _ImageSegmenter) + + def test_create_from_options_fails_with_invalid_model_path(self): + # Invalid empty model path. + with self.assertRaisesRegex( + ValueError, + r"ExternalFile must specify at least one of 'file_content', " + r"'file_name' or 'file_descriptor_meta'."): + base_options = _BaseOptions(file_name='') + options = _ImageSegmenterOptions(base_options=base_options) + _ImageSegmenter.create_from_options(options) + + def test_create_from_options_succeeds_with_valid_model_content(self): + # Creates with options containing model content successfully. + with open(self.model_path, 'rb') as f: + base_options = _BaseOptions(file_content=f.read()) + options = _ImageSegmenterOptions(base_options=base_options) + segmenter = _ImageSegmenter.create_from_options(options) + self.assertIsInstance(segmenter, _ImageSegmenter) + + @parameterized.parameters( + (ModelFileType.FILE_NAME, 4), + (ModelFileType.FILE_CONTENT, 4)) + def succeeds_with_category_mask(self, model_file_type, max_results): + # Creates segmenter. + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(file_name=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_content = f.read() + base_options = _BaseOptions(file_content=model_content) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + options = _ImageSegmenterOptions(base_options=base_options, + output_type=_OutputType.CATEGORY_MASK) + segmenter = _ImageSegmenter.create_from_options(options) + + # Performs image segmentation on the input. + image_result = segmenter.segment(self.test_image) + + # Comparing results. + print(image_result) + + # Closes the segmenter explicitly when the segmenter is not used in + # a context. + segmenter.close() + + +if __name__ == '__main__': + absltest.main() diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index 7ff81861..ce59763d 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -36,3 +36,22 @@ py_library( "//mediapipe/tasks/python/vision/core:vision_task_running_mode", ], ) + +py_library( + name = "image_segmenter", + srcs = [ + "image_segmenter.py", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/python:packet_creator", + "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_py_pb2", + "//mediapipe/tasks/python/components:segmenter_options", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/core:optional_dependencies", + "//mediapipe/tasks/python/core:task_info", + "//mediapipe/tasks/python/vision/core:base_vision_task_api", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py new file mode 100644 index 00000000..ea40d85c --- /dev/null +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -0,0 +1,205 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MediaPipe image segmenter task.""" + +import dataclasses +from typing import Callable, List, Mapping, Optional + +from mediapipe.python import packet_creator +from mediapipe.python import packet_getter +from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import packet as packet_module +from mediapipe.python._framework_bindings import task_runner as task_runner_module +from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_options_pb2 +from mediapipe.tasks.python.components import segmenter_options +from mediapipe.tasks.python.core import base_options as base_options_module +from mediapipe.tasks.python.core import task_info as task_info_module +from mediapipe.tasks.python.core.optional_dependencies import doc_controls +from mediapipe.tasks.python.vision.core import base_vision_task_api +from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module + +_BaseOptions = base_options_module.BaseOptions +_ImageSegmenterOptionsProto = image_segmenter_options_pb2.ImageSegmenterOptions +_SegmenterOptions = segmenter_options.SegmenterOptions +_RunningMode = running_mode_module.VisionTaskRunningMode +_TaskInfo = task_info_module.TaskInfo +_TaskRunner = task_runner_module.TaskRunner + +_SEGMENTATION_OUT_STREAM_NAME = 'segmented_masks' +_SEGMENTATION_TAG = 'SEGMENTATION' +_GROUPED_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' +_IMAGE_IN_STREAM_NAME = 'image_in' +_IMAGE_OUT_STREAM_NAME = 'image_out' +_IMAGE_TAG = 'IMAGE' +_TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph' + + +@dataclasses.dataclass +class ImageSegmenterOptions: + """Options for the image segmenter task. + + Attributes: + base_options: Base options for the image segmenter task. + running_mode: The running mode of the task. Default to the image mode. + Image segmenter task has three running modes: + 1) The image mode for detecting objects on single image inputs. + 2) The video mode for detecting objects on the decoded frames of a video. + 3) The live stream mode for detecting objects on a live stream of input + data, such as from camera. + output_type: Optional output mask type. + activation: Activation function to apply to input tensor. + result_callback: The user-defined result callback for processing live stream + data. The result callback should only be specified when the running mode + is set to the live stream mode. + """ + base_options: _BaseOptions + running_mode: _RunningMode = _RunningMode.IMAGE + output_type: Optional[segmenter_options.OutputType] = segmenter_options.OutputType.CATEGORY_MASK + activation: Optional[segmenter_options.Activation] = segmenter_options.Activation.NONE + result_callback: Optional[ + Callable[[List[image_module.Image], image_module.Image, int], + None]] = None + + @doc_controls.do_not_generate_docs + def to_pb2(self) -> _ImageSegmenterOptionsProto: + """Generates an ImageSegmenterOptions protobuf object.""" + base_options_proto = self.base_options.to_pb2() + base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True + + segmenter_options = _SegmenterOptions( + output_type=self.output_type, + activation=self.activation + ) + + return _ImageSegmenterOptionsProto( + base_options=base_options_proto, + segmenter_options=segmenter_options.to_pb2() + ) + + +class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): + """Class that performs image segmentation on images.""" + + @classmethod + def create_from_model_path(cls, model_path: str) -> 'ImageSegmenter': + """Creates an `ImageSegmenter` object from a TensorFlow Lite model and the default `ImageSegmenterOptions`. + + Note that the created `ImageSegmenter` instance is in image mode, for + performing image segmentation on single image inputs. + + Args: + model_path: Path to the model. + + Returns: + `ImageSegmenter` object that's created from the model file and the default + `ImageSegmenterOptions`. + + Raises: + ValueError: If failed to create `ImageSegmenter` object from the provided + file such as invalid file path. + RuntimeError: If other types of error occurred. + """ + base_options = _BaseOptions(file_name=model_path) + options = ImageSegmenterOptions( + base_options=base_options, running_mode=_RunningMode.IMAGE) + return cls.create_from_options(options) + + @classmethod + def create_from_options(cls, + options: ImageSegmenterOptions) -> 'ImageSegmenter': + """Creates the `ImageSegmenter` object from image segmenter options. + + Args: + options: Options for the image segmenter task. + + Returns: + `ImageSegmenter` object that's created from `options`. + + Raises: + ValueError: If failed to create `ImageSegmenter` object from + `ImageSegmenterOptions` such as missing the model. + RuntimeError: If other types of error occurred. + """ + + def packets_callback(output_packets: Mapping[str, packet_module.Packet]): + if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): + return + segmentation_result = packet_getter.get_proto_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) + timestamp = output_packets[_IMAGE_OUT_STREAM_NAME].timestamp + options.result_callback(segmentation_result, image, timestamp) + + task_info = _TaskInfo( + task_graph=_TASK_GRAPH_NAME, + input_streams=[':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME])], + output_streams=[ + ':'.join([_SEGMENTATION_TAG, _SEGMENTATION_OUT_STREAM_NAME]), + ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]) + ], + task_options=options) + return cls( + task_info.generate_graph_config( + enable_flow_limiting=options.running_mode == + _RunningMode.LIVE_STREAM), options.running_mode, + packets_callback if options.result_callback else None) + + # TODO: Create an Image class for MediaPipe Tasks. + def segment(self, + image: image_module.Image) -> List[image_module.Image]: + """Performs the actual segmentation task on the provided MediaPipe Image. + + Args: + image: MediaPipe Image. + + Returns: + A segmentation result object that contains a list of segmentation masks + as images. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If object detection failed to run. + """ + output_packets = self._process_image_data( + {_IMAGE_IN_STREAM_NAME: packet_creator.create_image(image)}) + segmentation_result = packet_getter.get_proto_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + return segmentation_result + + # def segment_async(self, image: image_module.Image, timestamp_ms: int) -> None: + # """Sends live image data (an Image with a unique timestamp) to perform image segmentation. + # + # This method will return immediately after the input image is accepted. The + # results will be available via the `result_callback` provided in the + # `ImageSegmenterOptions`. The `segment_async` method is designed to process + # live stream data such as camera input. To lower the overall latency, image + # segmenter may drop the input images if needed. In other words, it's not + # guaranteed to have output per input image. The `result_callback` provides: + # - A segmentation result object that contains a list of segmentation masks + # as images. + # - The input image that the image segmenter runs on. + # - The input timestamp in milliseconds. + # + # Args: + # image: MediaPipe Image. + # timestamp_ms: The timestamp of the input image in milliseconds. + # + # Raises: + # ValueError: If the current input timestamp is smaller than what the object + # detector has already processed. + # """ + # self._send_live_stream_data({ + # _IMAGE_IN_STREAM_NAME: + # packet_creator.create_image(image).at(timestamp_ms) + # }) From 500ad5a7f0263232bfa864968befce3ced6be010 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 21 Sep 2022 03:43:18 -0700 Subject: [PATCH 02/15] Updated some files --- mediapipe/python/BUILD | 1 - mediapipe/tasks/python/test/vision/image_segmenter_test.py | 5 +++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mediapipe/python/BUILD b/mediapipe/python/BUILD index 331ee836..55c0a22e 100644 --- a/mediapipe/python/BUILD +++ b/mediapipe/python/BUILD @@ -86,7 +86,6 @@ cc_library( name = "builtin_task_graphs", deps = [ "//mediapipe/tasks/cc/vision/object_detector:object_detector_graph", - "//mediapipe/tasks/cc/vision/image_classification:image_classifier_graph", "//mediapipe/tasks/cc/vision/image_segmenter:image_segmenter_graph", ], ) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 704b0fc5..6194f766 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -104,10 +104,11 @@ class ImageSegmenterTest(parameterized.TestCase): segmenter = _ImageSegmenter.create_from_options(options) # Performs image segmentation on the input. - image_result = segmenter.segment(self.test_image) + category_masks = segmenter.segment(self.test_image) # Comparing results. - print(image_result) + print(len(category_masks)) + s # Closes the segmenter explicitly when the segmenter is not used in # a context. From 660a88b7ead9c4aacd806e6fd7751db87cd428c1 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 21 Sep 2022 04:06:12 -0700 Subject: [PATCH 03/15] Code cleanup --- .../test/vision/image_segmenter_test.py | 9 ++-- .../tasks/python/vision/image_segmenter.py | 41 ++----------------- 2 files changed, 7 insertions(+), 43 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 6194f766..8a91766d 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -17,7 +17,6 @@ import enum from absl.testing import absltest from absl.testing import parameterized -import numpy as np from mediapipe.python._framework_bindings import image as image_module from mediapipe.tasks.python.components import segmenter_options @@ -30,6 +29,7 @@ _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _OutputType = segmenter_options.OutputType _Activation = segmenter_options.Activation +_SegmenterOptions = segmenter_options.SegmenterOptions _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions _RUNNING_MODE = running_mode_module.VisionTaskRunningMode @@ -99,17 +99,14 @@ class ImageSegmenterTest(parameterized.TestCase): # Should never happen raise ValueError('model_file_type is invalid.') + segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions(base_options=base_options, - output_type=_OutputType.CATEGORY_MASK) + segmenter_options=segmenter_options) segmenter = _ImageSegmenter.create_from_options(options) # Performs image segmentation on the input. category_masks = segmenter.segment(self.test_image) - # Comparing results. - print(len(category_masks)) - s - # Closes the segmenter explicitly when the segmenter is not used in # a context. segmenter.close() diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index ea40d85c..47f09647 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -57,16 +57,14 @@ class ImageSegmenterOptions: 2) The video mode for detecting objects on the decoded frames of a video. 3) The live stream mode for detecting objects on a live stream of input data, such as from camera. - output_type: Optional output mask type. - activation: Activation function to apply to input tensor. + segmenter_options: Options for the image segmenter task. result_callback: The user-defined result callback for processing live stream data. The result callback should only be specified when the running mode is set to the live stream mode. """ base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE - output_type: Optional[segmenter_options.OutputType] = segmenter_options.OutputType.CATEGORY_MASK - activation: Optional[segmenter_options.Activation] = segmenter_options.Activation.NONE + segmenter_options: _SegmenterOptions = _SegmenterOptions() result_callback: Optional[ Callable[[List[image_module.Image], image_module.Image, int], None]] = None @@ -76,15 +74,11 @@ class ImageSegmenterOptions: """Generates an ImageSegmenterOptions protobuf object.""" base_options_proto = self.base_options.to_pb2() base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True - - segmenter_options = _SegmenterOptions( - output_type=self.output_type, - activation=self.activation - ) + segmenter_options_proto = self.segmenter_options.to_pb2() return _ImageSegmenterOptionsProto( base_options=base_options_proto, - segmenter_options=segmenter_options.to_pb2() + segmenter_options=segmenter_options_proto ) @@ -176,30 +170,3 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): segmentation_result = packet_getter.get_proto_list( output_packets[_SEGMENTATION_OUT_STREAM_NAME]) return segmentation_result - - # def segment_async(self, image: image_module.Image, timestamp_ms: int) -> None: - # """Sends live image data (an Image with a unique timestamp) to perform image segmentation. - # - # This method will return immediately after the input image is accepted. The - # results will be available via the `result_callback` provided in the - # `ImageSegmenterOptions`. The `segment_async` method is designed to process - # live stream data such as camera input. To lower the overall latency, image - # segmenter may drop the input images if needed. In other words, it's not - # guaranteed to have output per input image. The `result_callback` provides: - # - A segmentation result object that contains a list of segmentation masks - # as images. - # - The input image that the image segmenter runs on. - # - The input timestamp in milliseconds. - # - # Args: - # image: MediaPipe Image. - # timestamp_ms: The timestamp of the input image in milliseconds. - # - # Raises: - # ValueError: If the current input timestamp is smaller than what the object - # detector has already processed. - # """ - # self._send_live_stream_data({ - # _IMAGE_IN_STREAM_NAME: - # packet_creator.create_image(image).at(timestamp_ms) - # }) From e028b24c42c8df304551ac6c55ff7706f70855dd Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 21 Sep 2022 04:34:19 -0700 Subject: [PATCH 04/15] Updated values for some constants --- mediapipe/tasks/python/vision/image_segmenter.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 47f09647..217867d4 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -36,9 +36,8 @@ _RunningMode = running_mode_module.VisionTaskRunningMode _TaskInfo = task_info_module.TaskInfo _TaskRunner = task_runner_module.TaskRunner -_SEGMENTATION_OUT_STREAM_NAME = 'segmented_masks' -_SEGMENTATION_TAG = 'SEGMENTATION' -_GROUPED_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' +_SEGMENTATION_OUT_STREAM_NAME = 'segmented_mask_out' +_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' _IMAGE_IN_STREAM_NAME = 'image_in' _IMAGE_OUT_STREAM_NAME = 'image_out' _IMAGE_TAG = 'IMAGE' From d25626ff63fee6352c6a4f7d15159875c62c5541 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Sun, 25 Sep 2022 09:16:13 -0700 Subject: [PATCH 05/15] Added Image Segmenter implementation and tests --- mediapipe/python/BUILD | 3 +- mediapipe/tasks/python/test/vision/BUILD | 1 + .../test/vision/image_segmenter_test.py | 90 ++++++++++++++++++- .../tasks/python/vision/image_segmenter.py | 4 +- 4 files changed, 92 insertions(+), 6 deletions(-) diff --git a/mediapipe/python/BUILD b/mediapipe/python/BUILD index 55c0a22e..98e7f80c 100644 --- a/mediapipe/python/BUILD +++ b/mediapipe/python/BUILD @@ -35,7 +35,7 @@ pybind_extension( }), module_name = "_framework_bindings", deps = [ - ":builtin_calculators", + #":builtin_calculators", ":builtin_task_graphs", "//mediapipe/python/pybind:calculator_graph", "//mediapipe/python/pybind:image", @@ -85,6 +85,7 @@ cc_library( cc_library( name = "builtin_task_graphs", deps = [ + "//mediapipe/calculators/core:flow_limiter_calculator", "//mediapipe/tasks/cc/vision/object_detector:object_detector_graph", "//mediapipe/tasks/cc/vision/image_segmenter:image_segmenter_graph", ], diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 403e00a3..bec75923 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -47,6 +47,7 @@ py_test( ], deps = [ # build rule placeholder: numpy dep, + # build rule placeholder: cv2 dep, "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/test:test_util", "//mediapipe/tasks/python/components:segmenter_options", diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 8a91766d..c3763577 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -14,11 +14,14 @@ """Tests for image segmenter.""" import enum +import numpy as np +import cv2 from absl.testing import absltest from absl.testing import parameterized from mediapipe.python._framework_bindings import image as image_module +from mediapipe.python._framework_bindings import image_frame as image_frame_module from mediapipe.tasks.python.components import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.test import test_util @@ -27,6 +30,7 @@ from mediapipe.tasks.python.vision.core import vision_task_running_mode as runni _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image +_ImageFormat = image_frame_module.ImageFormat _OutputType = segmenter_options.OutputType _Activation = segmenter_options.Activation _SegmenterOptions = segmenter_options.SegmenterOptions @@ -41,6 +45,13 @@ _MASK_MAGNIFICATION_FACTOR = 10 _MATCH_PIXELS_THRESHOLD = 0.01 +def _iou(ground_truth, prediction): + intersection = np.logical_and(ground_truth, prediction) + union = np.logical_or(ground_truth, prediction) + iou = np.sum(intersection) / np.sum(union) + return iou + + class ModelFileType(enum.Enum): FILE_CONTENT = 1 FILE_NAME = 2 @@ -52,6 +63,7 @@ class ImageSegmenterTest(parameterized.TestCase): super().setUp() self.test_image = test_util.read_test_image( test_util.get_test_data_path(_IMAGE_FILE)) + self.test_seg_path = test_util.get_test_data_path(_SEGMENTATION_FILE) self.model_path = test_util.get_test_data_path(_MODEL_FILE) def test_create_from_file_succeeds_with_valid_model_path(self): @@ -85,9 +97,9 @@ class ImageSegmenterTest(parameterized.TestCase): self.assertIsInstance(segmenter, _ImageSegmenter) @parameterized.parameters( - (ModelFileType.FILE_NAME, 4), - (ModelFileType.FILE_CONTENT, 4)) - def succeeds_with_category_mask(self, model_file_type, max_results): + (ModelFileType.FILE_NAME,), + (ModelFileType.FILE_CONTENT,)) + def test_succeeds_with_category_mask(self, model_file_type): # Creates segmenter. if model_file_type is ModelFileType.FILE_NAME: base_options = _BaseOptions(file_name=self.model_path) @@ -106,6 +118,78 @@ class ImageSegmenterTest(parameterized.TestCase): # Performs image segmentation on the input. category_masks = segmenter.segment(self.test_image) + self.assertEqual(len(category_masks), 1) + result_pixels = category_masks[0].numpy_view().flatten() + + # Check if data type of `category_masks` is correct. + self.assertEqual(result_pixels.dtype, np.uint8) + + # Loads ground truth segmentation file. + image_data = cv2.imread(self.test_seg_path, cv2.IMREAD_GRAYSCALE) + gt_segmentation = _Image(_ImageFormat.GRAY8, image_data) + gt_segmentation_array = gt_segmentation.numpy_view() + gt_segmentation_shape = gt_segmentation_array.shape + num_pixels = gt_segmentation_shape[0] * gt_segmentation_shape[1] + ground_truth_pixels = gt_segmentation_array.flatten() + + self.assertEqual( + len(result_pixels), len(ground_truth_pixels), + 'Segmentation mask size does not match the ground truth mask size.') + + inconsistent_pixels = 0 + + for index in range(num_pixels): + inconsistent_pixels += ( + result_pixels[index] * _MASK_MAGNIFICATION_FACTOR != + ground_truth_pixels[index]) + + self.assertLessEqual( + inconsistent_pixels / num_pixels, _MATCH_PIXELS_THRESHOLD, + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MATCH_PIXELS_THRESHOLD}.') + + # Closes the segmenter explicitly when the segmenter is not used in + # a context. + segmenter.close() + + def test_succeeds_with_confidence_mask(self): + # Creates segmenter. + base_options = _BaseOptions(file_name=self.model_path) + + # Run segmentation on the model in CATEGORY_MASK mode. + segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) + options = _ImageSegmenterOptions(base_options=base_options, + segmenter_options=segmenter_options) + segmenter = _ImageSegmenter.create_from_options(options) + category_masks = segmenter.segment(self.test_image) + category_mask = category_masks[0].numpy_view() + + # Run segmentation on the model in CONFIDENCE_MASK mode. + segmenter_options = _SegmenterOptions( + output_type=_OutputType.CONFIDENCE_MASK, + activation=_Activation.SOFTMAX) + options = _ImageSegmenterOptions(base_options=base_options, + segmenter_options=segmenter_options) + segmenter = _ImageSegmenter.create_from_options(options) + confidence_masks = segmenter.segment(self.test_image) + + # Check if confidence mask shape is correct. + self.assertEqual( + len(confidence_masks), 21, + 'Number of confidence masks must match with number of categories.') + + # Gather the confidence masks in a single array `confidence_mask_array`. + confidence_mask_array = np.array( + [confidence_mask.numpy_view() for confidence_mask in confidence_masks]) + + # Check if data type of `confidence_masks` are correct. + self.assertEqual(confidence_mask_array.dtype, np.float32) + + # Compute the category mask from the created confidence mask. + calculated_category_mask = np.argmax(confidence_mask_array, axis=0) + self.assertListEqual( + calculated_category_mask.tolist(), category_mask.tolist(), + 'Confidence mask does not match with the category mask.') # Closes the segmenter explicitly when the segmenter is not used in # a context. diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 217867d4..b643b1cb 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -128,7 +128,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): def packets_callback(output_packets: Mapping[str, packet_module.Packet]): if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): return - segmentation_result = packet_getter.get_proto_list( + segmentation_result = packet_getter.get_image_list( output_packets[_SEGMENTATION_OUT_STREAM_NAME]) image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) timestamp = output_packets[_IMAGE_OUT_STREAM_NAME].timestamp @@ -166,6 +166,6 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): """ output_packets = self._process_image_data( {_IMAGE_IN_STREAM_NAME: packet_creator.create_image(image)}) - segmentation_result = packet_getter.get_proto_list( + segmentation_result = packet_getter.get_image_list( output_packets[_SEGMENTATION_OUT_STREAM_NAME]) return segmentation_result From 63e0c042535c8924862fd0b8c82c96223ea3e865 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 29 Sep 2022 02:03:07 -0700 Subject: [PATCH 06/15] Updated BUILD and tests --- mediapipe/tasks/python/test/vision/BUILD | 6 ++---- mediapipe/tasks/python/test/vision/image_segmenter_test.py | 7 ------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index bec75923..6baca885 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -26,7 +26,7 @@ py_test( "//mediapipe/tasks/testdata/vision:test_models", ], deps = [ - # build rule placeholder: numpy dep, + "//mediapipe/python:_framework_bindings", "//mediapipe/tasks/python/components/containers:bounding_box", "//mediapipe/tasks/python/components/containers:category", "//mediapipe/tasks/python/components/containers:detections", @@ -34,7 +34,6 @@ py_test( "//mediapipe/tasks/python/test:test_util", "//mediapipe/tasks/python/vision:object_detector", "//mediapipe/tasks/python/vision/core:vision_task_running_mode", - "@absl_py//absl/testing:parameterized", ], ) @@ -46,8 +45,7 @@ py_test( "//mediapipe/tasks/testdata/vision:test_models", ], deps = [ - # build rule placeholder: numpy dep, - # build rule placeholder: cv2 dep, + "//mediapipe/python:_framework_bindings", "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/test:test_util", "//mediapipe/tasks/python/components:segmenter_options", diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index c3763577..72be676d 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -45,13 +45,6 @@ _MASK_MAGNIFICATION_FACTOR = 10 _MATCH_PIXELS_THRESHOLD = 0.01 -def _iou(ground_truth, prediction): - intersection = np.logical_and(ground_truth, prediction) - union = np.logical_or(ground_truth, prediction) - iou = np.sum(intersection) / np.sum(union) - return iou - - class ModelFileType(enum.Enum): FILE_CONTENT = 1 FILE_NAME = 2 From bef2f6ccedaca0ac28f45602a7676fa794090f57 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 29 Sep 2022 03:40:56 -0700 Subject: [PATCH 07/15] Updated implementation and tests --- mediapipe/tasks/python/components/BUILD | 9 ------ mediapipe/tasks/python/components/proto/BUILD | 28 +++++++++++++++++++ .../{ => proto}/segmenter_options.py | 2 +- mediapipe/tasks/python/test/vision/BUILD | 2 +- .../test/vision/image_segmenter_test.py | 14 +++++----- mediapipe/tasks/python/vision/BUILD | 2 +- .../tasks/python/vision/image_segmenter.py | 4 +-- 7 files changed, 40 insertions(+), 21 deletions(-) create mode 100644 mediapipe/tasks/python/components/proto/BUILD rename mediapipe/tasks/python/components/{ => proto}/segmenter_options.py (97%) diff --git a/mediapipe/tasks/python/components/BUILD b/mediapipe/tasks/python/components/BUILD index eb8714a9..00fd4061 100644 --- a/mediapipe/tasks/python/components/BUILD +++ b/mediapipe/tasks/python/components/BUILD @@ -17,12 +17,3 @@ package(default_visibility = ["//mediapipe/tasks:internal"]) licenses(["notice"]) - -py_library( - name = "segmenter_options", - srcs = ["segmenter_options.py"], - deps = [ - "//mediapipe/tasks/cc/components:segmenter_options_py_pb2", - "//mediapipe/tasks/python/core:optional_dependencies", - ], -) diff --git a/mediapipe/tasks/python/components/proto/BUILD b/mediapipe/tasks/python/components/proto/BUILD new file mode 100644 index 00000000..a58f77e6 --- /dev/null +++ b/mediapipe/tasks/python/components/proto/BUILD @@ -0,0 +1,28 @@ +# Copyright 2022 The MediaPipe Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Placeholder for internal Python strict library compatibility macro. + +package(default_visibility = ["//mediapipe/tasks:internal"]) + +licenses(["notice"]) + +py_library( + name = "segmenter_options", + srcs = ["segmenter_options.py"], + deps = [ + "//mediapipe/tasks/cc/components/proto:segmenter_options_py_pb2", + "//mediapipe/tasks/python/core:optional_dependencies", + ], +) diff --git a/mediapipe/tasks/python/components/segmenter_options.py b/mediapipe/tasks/python/components/proto/segmenter_options.py similarity index 97% rename from mediapipe/tasks/python/components/segmenter_options.py rename to mediapipe/tasks/python/components/proto/segmenter_options.py index 5b94a256..dcf34cc3 100644 --- a/mediapipe/tasks/python/components/segmenter_options.py +++ b/mediapipe/tasks/python/components/proto/segmenter_options.py @@ -17,7 +17,7 @@ import dataclasses import enum from typing import Any, Optional -from mediapipe.tasks.cc.components import segmenter_options_pb2 +from mediapipe.tasks.cc.components.proto import segmenter_options_pb2 from mediapipe.tasks.python.core.optional_dependencies import doc_controls _SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 6baca885..ea6608fc 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -48,7 +48,7 @@ py_test( "//mediapipe/python:_framework_bindings", "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/test:test_util", - "//mediapipe/tasks/python/components:segmenter_options", + "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/vision:image_segmenter", "//mediapipe/tasks/python/vision/core:vision_task_running_mode", "@absl_py//absl/testing:parameterized", diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 72be676d..9ebd46bc 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -22,7 +22,7 @@ from absl.testing import parameterized from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import image_frame as image_frame_module -from mediapipe.tasks.python.components import segmenter_options +from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.test import test_util from mediapipe.tasks.python.vision import image_segmenter @@ -66,7 +66,7 @@ class ImageSegmenterTest(parameterized.TestCase): def test_create_from_options_succeeds_with_valid_model_path(self): # Creates with options containing model file successfully. - base_options = _BaseOptions(file_name=self.model_path) + base_options = _BaseOptions(model_asset_path=self.model_path) options = _ImageSegmenterOptions(base_options=base_options) with _ImageSegmenter.create_from_options(options) as segmenter: self.assertIsInstance(segmenter, _ImageSegmenter) @@ -77,14 +77,14 @@ class ImageSegmenterTest(parameterized.TestCase): ValueError, r"ExternalFile must specify at least one of 'file_content', " r"'file_name' or 'file_descriptor_meta'."): - base_options = _BaseOptions(file_name='') + base_options = _BaseOptions(model_asset_path='') options = _ImageSegmenterOptions(base_options=base_options) _ImageSegmenter.create_from_options(options) def test_create_from_options_succeeds_with_valid_model_content(self): # Creates with options containing model content successfully. with open(self.model_path, 'rb') as f: - base_options = _BaseOptions(file_content=f.read()) + base_options = _BaseOptions(model_asset_buffer=f.read()) options = _ImageSegmenterOptions(base_options=base_options) segmenter = _ImageSegmenter.create_from_options(options) self.assertIsInstance(segmenter, _ImageSegmenter) @@ -95,11 +95,11 @@ class ImageSegmenterTest(parameterized.TestCase): def test_succeeds_with_category_mask(self, model_file_type): # Creates segmenter. if model_file_type is ModelFileType.FILE_NAME: - base_options = _BaseOptions(file_name=self.model_path) + base_options = _BaseOptions(model_asset_path=self.model_path) elif model_file_type is ModelFileType.FILE_CONTENT: with open(self.model_path, 'rb') as f: model_content = f.read() - base_options = _BaseOptions(file_content=model_content) + base_options = _BaseOptions(model_asset_buffer=model_content) else: # Should never happen raise ValueError('model_file_type is invalid.') @@ -147,7 +147,7 @@ class ImageSegmenterTest(parameterized.TestCase): def test_succeeds_with_confidence_mask(self): # Creates segmenter. - base_options = _BaseOptions(file_name=self.model_path) + base_options = _BaseOptions(model_asset_path=self.model_path) # Run segmentation on the model in CATEGORY_MASK mode. segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index ce59763d..3875ea5d 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -47,7 +47,7 @@ py_library( "//mediapipe/python:packet_creator", "//mediapipe/python:packet_getter", "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_py_pb2", - "//mediapipe/tasks/python/components:segmenter_options", + "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/core:optional_dependencies", "//mediapipe/tasks/python/core:task_info", diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index b643b1cb..060a6779 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -22,7 +22,7 @@ from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import packet as packet_module from mediapipe.python._framework_bindings import task_runner as task_runner_module from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_options_pb2 -from mediapipe.tasks.python.components import segmenter_options +from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.core import task_info as task_info_module from mediapipe.tasks.python.core.optional_dependencies import doc_controls @@ -103,7 +103,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): file such as invalid file path. RuntimeError: If other types of error occurred. """ - base_options = _BaseOptions(file_name=model_path) + base_options = _BaseOptions(model_asset_path=model_path) options = ImageSegmenterOptions( base_options=base_options, running_mode=_RunningMode.IMAGE) return cls.create_from_options(options) From f84e0bc1c61a0b4de6a296d08bba6934bcf2f18d Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Oct 2022 04:24:12 -0700 Subject: [PATCH 08/15] Revised API implementation and added more tests for segment_for_video and segment_async --- mediapipe/tasks/python/test/vision/BUILD | 2 +- .../test/vision/image_segmenter_test.py | 233 +++++++++++++++--- .../tasks/python/vision/image_segmenter.py | 82 +++++- 3 files changed, 277 insertions(+), 40 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 51a9c514..63fc56b4 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -47,7 +47,7 @@ py_test( deps = [ "//mediapipe/python:_framework_bindings", "//mediapipe/tasks/python/core:base_options", - "//mediapipe/tasks/python/test:test_util", + "//mediapipe/tasks/python/test:test_utils", "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/vision:image_segmenter", "//mediapipe/tasks/python/vision/core:vision_task_running_mode", diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 9ebd46bc..a97aed10 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -16,6 +16,8 @@ import enum import numpy as np import cv2 +from typing import List +from unittest import mock from absl.testing import absltest from absl.testing import parameterized @@ -24,7 +26,7 @@ from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import image_frame as image_frame_module from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module -from mediapipe.tasks.python.test import test_util +from mediapipe.tasks.python.test import test_utils from mediapipe.tasks.python.vision import image_segmenter from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module @@ -42,7 +44,22 @@ _MODEL_FILE = 'deeplabv3.tflite' _IMAGE_FILE = 'segmentation_input_rotation0.jpg' _SEGMENTATION_FILE = 'segmentation_golden_rotation0.png' _MASK_MAGNIFICATION_FACTOR = 10 -_MATCH_PIXELS_THRESHOLD = 0.01 +_MASK_SIMILARITY_THRESHOLD = 0.98 + + +def _similar_to_uint8_mask(actual_mask, expected_mask): + actual_mask_pixels = actual_mask.numpy_view().flatten() + expected_mask_pixels = expected_mask.numpy_view().flatten() + + consistent_pixels = 0 + num_pixels = len(expected_mask_pixels) + + for index in range(num_pixels): + consistent_pixels += ( + actual_mask_pixels[index] * _MASK_MAGNIFICATION_FACTOR == + expected_mask_pixels[index]) + + return consistent_pixels / num_pixels >= _MASK_SIMILARITY_THRESHOLD class ModelFileType(enum.Enum): @@ -54,10 +71,14 @@ class ImageSegmenterTest(parameterized.TestCase): def setUp(self): super().setUp() - self.test_image = test_util.read_test_image( - test_util.get_test_data_path(_IMAGE_FILE)) - self.test_seg_path = test_util.get_test_data_path(_SEGMENTATION_FILE) - self.model_path = test_util.get_test_data_path(_MODEL_FILE) + # Load the test input image. + self.test_image = _Image.create_from_file( + test_utils.get_test_data_path(_IMAGE_FILE)) + # Loads ground truth segmentation file. + gt_segmentation_data = cv2.imread( + test_utils.get_test_data_path(_SEGMENTATION_FILE), cv2.IMREAD_GRAYSCALE) + self.test_seg_image = _Image(_ImageFormat.GRAY8, gt_segmentation_data) + self.model_path = test_utils.get_test_data_path(_MODEL_FILE) def test_create_from_file_succeeds_with_valid_model_path(self): # Creates with default option and valid model file successfully. @@ -76,7 +97,7 @@ class ImageSegmenterTest(parameterized.TestCase): with self.assertRaisesRegex( ValueError, r"ExternalFile must specify at least one of 'file_content', " - r"'file_name' or 'file_descriptor_meta'."): + r"'file_name', 'file_pointer_meta' or 'file_descriptor_meta'."): base_options = _BaseOptions(model_asset_path='') options = _ImageSegmenterOptions(base_options=base_options) _ImageSegmenter.create_from_options(options) @@ -112,34 +133,16 @@ class ImageSegmenterTest(parameterized.TestCase): # Performs image segmentation on the input. category_masks = segmenter.segment(self.test_image) self.assertEqual(len(category_masks), 1) - result_pixels = category_masks[0].numpy_view().flatten() + category_mask = category_masks[0] + result_pixels = category_mask.numpy_view().flatten() - # Check if data type of `category_masks` is correct. + # Check if data type of `category_mask` is correct. self.assertEqual(result_pixels.dtype, np.uint8) - # Loads ground truth segmentation file. - image_data = cv2.imread(self.test_seg_path, cv2.IMREAD_GRAYSCALE) - gt_segmentation = _Image(_ImageFormat.GRAY8, image_data) - gt_segmentation_array = gt_segmentation.numpy_view() - gt_segmentation_shape = gt_segmentation_array.shape - num_pixels = gt_segmentation_shape[0] * gt_segmentation_shape[1] - ground_truth_pixels = gt_segmentation_array.flatten() - - self.assertEqual( - len(result_pixels), len(ground_truth_pixels), - 'Segmentation mask size does not match the ground truth mask size.') - - inconsistent_pixels = 0 - - for index in range(num_pixels): - inconsistent_pixels += ( - result_pixels[index] * _MASK_MAGNIFICATION_FACTOR != - ground_truth_pixels[index]) - - self.assertLessEqual( - inconsistent_pixels / num_pixels, _MATCH_PIXELS_THRESHOLD, + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), f'Number of pixels in the candidate mask differing from that of the ' - f'ground truth mask exceeds {_MATCH_PIXELS_THRESHOLD}.') + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') # Closes the segmenter explicitly when the segmenter is not used in # a context. @@ -188,6 +191,174 @@ class ImageSegmenterTest(parameterized.TestCase): # a context. segmenter.close() + @parameterized.parameters( + (ModelFileType.FILE_NAME,), + (ModelFileType.FILE_CONTENT,)) + def test_segment_in_context(self, model_file_type): + if model_file_type is ModelFileType.FILE_NAME: + base_options = _BaseOptions(model_asset_path=self.model_path) + elif model_file_type is ModelFileType.FILE_CONTENT: + with open(self.model_path, 'rb') as f: + model_contents = f.read() + base_options = _BaseOptions(model_asset_buffer=model_contents) + else: + # Should never happen + raise ValueError('model_file_type is invalid.') + + segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) + options = _ImageSegmenterOptions(base_options=base_options, + segmenter_options=segmenter_options) + with _ImageSegmenter.create_from_options(options) as segmenter: + # Performs image segmentation on the input. + category_masks = segmenter.segment(self.test_image) + self.assertEqual(len(category_masks), 1) + + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + + def test_missing_result_callback(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM) + with self.assertRaisesRegex(ValueError, + r'result callback must be provided'): + with _ImageSegmenter.create_from_options(options) as unused_segmenter: + pass + + @parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO)) + def test_illegal_result_callback(self, running_mode): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=running_mode, + result_callback=mock.MagicMock()) + with self.assertRaisesRegex(ValueError, + r'result callback should not be provided'): + with _ImageSegmenter.create_from_options(options) as unused_segmenter: + pass + + def test_calling_segment_for_video_in_image_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + segmenter.segment_for_video(self.test_image, 0) + + def test_calling_segment_async_in_image_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.IMAGE) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + segmenter.segment_async(self.test_image, 0) + + def test_calling_segment_in_video_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + segmenter.segment(self.test_image) + + def test_calling_segment_async_in_video_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the live stream mode'): + segmenter.segment_async(self.test_image, 0) + + def test_detect_for_video_with_out_of_order_timestamp(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + unused_result = segmenter.segment_for_video(self.test_image, 1) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + segmenter.segment_for_video(self.test_image, 0) + + def test_segment_for_video(self): + segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + segmenter_options=segmenter_options, + running_mode=_RUNNING_MODE.VIDEO) + with _ImageSegmenter.create_from_options(options) as segmenter: + for timestamp in range(0, 300, 30): + category_masks = segmenter.segment_for_video(self.test_image, timestamp) + self.assertEqual(len(category_masks), 1) + self.assertTrue( + _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + + def test_calling_segment_in_live_stream_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the image mode'): + segmenter.segment(self.test_image) + + def test_calling_segment_for_video_in_live_stream_mode(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + with self.assertRaisesRegex(ValueError, + r'not initialized with the video mode'): + segmenter.segment_for_video(self.test_image, 0) + + def test_segment_async_calls_with_illegal_timestamp(self): + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=mock.MagicMock()) + with _ImageSegmenter.create_from_options(options) as segmenter: + segmenter.segment_async(self.test_image, 100) + with self.assertRaisesRegex( + ValueError, r'Input timestamp must be monotonically increasing'): + segmenter.segment_async(self.test_image, 0) + + def test_segment_async_calls(self): + observed_timestamp_ms = -1 + + def check_result(result: List[image_module.Image], + output_image: _Image, + timestamp_ms: int): + # Get the output category mask. + category_mask = result[0] + self.assertEqual(output_image.width, self.test_image.width) + self.assertEqual(output_image.height, self.test_image.height) + self.assertEqual(output_image.width, self.test_seg_image.width) + self.assertEqual(output_image.height, self.test_seg_image.height) + self.assertTrue( + _similar_to_uint8_mask(category_mask, self.test_seg_image), + f'Number of pixels in the candidate mask differing from that of the ' + f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + self.assertLess(observed_timestamp_ms, timestamp_ms) + self.observed_timestamp_ms = timestamp_ms + + segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + segmenter_options=segmenter_options, + running_mode=_RUNNING_MODE.LIVE_STREAM, + result_callback=check_result) + with _ImageSegmenter.create_from_options(options) as segmenter: + for timestamp in range(0, 300, 30): + segmenter.segment_async(self.test_image, timestamp) + if __name__ == '__main__': absltest.main() diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 060a6779..51f80292 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -42,6 +42,7 @@ _IMAGE_IN_STREAM_NAME = 'image_in' _IMAGE_OUT_STREAM_NAME = 'image_out' _IMAGE_TAG = 'IMAGE' _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph' +_MICRO_SECONDS_PER_MILLISECOND = 1000 @dataclasses.dataclass @@ -52,9 +53,9 @@ class ImageSegmenterOptions: base_options: Base options for the image segmenter task. running_mode: The running mode of the task. Default to the image mode. Image segmenter task has three running modes: - 1) The image mode for detecting objects on single image inputs. - 2) The video mode for detecting objects on the decoded frames of a video. - 3) The live stream mode for detecting objects on a live stream of input + 1) The image mode for segmenting objects on single image inputs. + 2) The video mode for segmenting objects on the decoded frames of a video. + 3) The live stream mode for segmenting objects on a live stream of input data, such as from camera. segmenter_options: Options for the image segmenter task. result_callback: The user-defined result callback for processing live stream @@ -86,7 +87,8 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): @classmethod def create_from_model_path(cls, model_path: str) -> 'ImageSegmenter': - """Creates an `ImageSegmenter` object from a TensorFlow Lite model and the default `ImageSegmenterOptions`. + """Creates an `ImageSegmenter` object from a TensorFlow Lite model and the + default `ImageSegmenterOptions`. Note that the created `ImageSegmenter` instance is in image mode, for performing image segmentation on single image inputs. @@ -131,8 +133,9 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): segmentation_result = packet_getter.get_image_list( output_packets[_SEGMENTATION_OUT_STREAM_NAME]) image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) - timestamp = output_packets[_IMAGE_OUT_STREAM_NAME].timestamp - options.result_callback(segmentation_result, image, timestamp) + timestamp = output_packets[_SEGMENTATION_OUT_STREAM_NAME].timestamp + options.result_callback(segmentation_result, image, + timestamp.value // _MICRO_SECONDS_PER_MILLISECOND) task_info = _TaskInfo( task_graph=_TASK_GRAPH_NAME, @@ -148,7 +151,6 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): _RunningMode.LIVE_STREAM), options.running_mode, packets_callback if options.result_callback else None) - # TODO: Create an Image class for MediaPipe Tasks. def segment(self, image: image_module.Image) -> List[image_module.Image]: """Performs the actual segmentation task on the provided MediaPipe Image. @@ -162,10 +164,74 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): Raises: ValueError: If any of the input arguments is invalid. - RuntimeError: If object detection failed to run. + RuntimeError: If image segmentation failed to run. """ output_packets = self._process_image_data( {_IMAGE_IN_STREAM_NAME: packet_creator.create_image(image)}) segmentation_result = packet_getter.get_image_list( output_packets[_SEGMENTATION_OUT_STREAM_NAME]) return segmentation_result + + def segment_for_video(self, image: image_module.Image, + timestamp_ms: int) -> List[image_module.Image]: + """Performs segmentation on the provided video frames. + + Only use this method when the ImageSegmenter is created with the video + running mode. It's required to provide the video frame's timestamp (in + milliseconds) along with the video frame. The input timestamps should be + monotonically increasing for adjacent calls of this method. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input video frame in milliseconds. + + Returns: + A segmentation result object that contains a list of segmentation masks + as images. + + Raises: + ValueError: If any of the input arguments is invalid. + RuntimeError: If image segmentation failed to run. + """ + output_packets = self._process_video_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) + segmentation_result = packet_getter.get_image_list( + output_packets[_SEGMENTATION_OUT_STREAM_NAME]) + return segmentation_result + + def segment_async(self, image: image_module.Image, timestamp_ms: int) -> None: + """Sends live image data (an Image with a unique timestamp) to perform + image segmentation. + + Only use this method when the ImageSegmenter is created with the live stream + running mode. The input timestamps should be monotonically increasing for + adjacent calls of this method. This method will return immediately after the + input image is accepted. The results will be available via the + `result_callback` provided in the `ImageSegmenterOptions`. The + `segment_async` method is designed to process live stream data such as + camera input. To lower the overall latency, image segmenter may drop the + input images if needed. In other words, it's not guaranteed to have output + per input image. + + The `result_callback` prvoides: + - A segmentation result object that contains a list of segmentation masks + as images. + - The input image that the image segmenter runs on. + - The input timestamp in milliseconds. + + Args: + image: MediaPipe Image. + timestamp_ms: The timestamp of the input image in milliseconds. + + Raises: + ValueError: If the current input timestamp is smaller than what the image + segmenter has already processed. + """ + self._send_live_stream_data({ + _IMAGE_IN_STREAM_NAME: + packet_creator.create_image(image).at( + timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND) + }) From 4932844410b0d2f7f18d7cd63b5b9f4915c7ea86 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Oct 2022 04:29:28 -0700 Subject: [PATCH 09/15] Reverted changes to BUILD --- mediapipe/python/BUILD | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mediapipe/python/BUILD b/mediapipe/python/BUILD index a329c44e..07ad9781 100644 --- a/mediapipe/python/BUILD +++ b/mediapipe/python/BUILD @@ -35,7 +35,7 @@ pybind_extension( }), module_name = "_framework_bindings", deps = [ - #":builtin_calculators", + ":builtin_calculators", ":builtin_task_graphs", "//mediapipe/python/pybind:calculator_graph", "//mediapipe/python/pybind:image", @@ -87,7 +87,6 @@ cc_library( cc_library( name = "builtin_task_graphs", deps = [ - "//mediapipe/calculators/core:flow_limiter_calculator", "//mediapipe/tasks/cc/vision/object_detector:object_detector_graph", "//mediapipe/tasks/cc/vision/image_segmenter:image_segmenter_graph", ], From 69f4daf9bc67c2d2ba03e9737a027cd09bf51b1e Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Oct 2022 04:31:35 -0700 Subject: [PATCH 10/15] Removed unused BUILD --- mediapipe/tasks/python/components/BUILD | 19 ------------------- 1 file changed, 19 deletions(-) delete mode 100644 mediapipe/tasks/python/components/BUILD diff --git a/mediapipe/tasks/python/components/BUILD b/mediapipe/tasks/python/components/BUILD deleted file mode 100644 index 00fd4061..00000000 --- a/mediapipe/tasks/python/components/BUILD +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright 2022 The MediaPipe Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Placeholder for internal Python strict library compatibility macro. - -package(default_visibility = ["//mediapipe/tasks:internal"]) - -licenses(["notice"]) From f166eb32e88df10ef6c56f63a6acaa5fac4dc0c3 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Oct 2022 04:34:17 -0700 Subject: [PATCH 11/15] Updated names of test case methods --- mediapipe/tasks/python/test/vision/image_segmenter_test.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index a97aed10..054c3e96 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -113,7 +113,7 @@ class ImageSegmenterTest(parameterized.TestCase): @parameterized.parameters( (ModelFileType.FILE_NAME,), (ModelFileType.FILE_CONTENT,)) - def test_succeeds_with_category_mask(self, model_file_type): + def test_segment_succeeds_with_category_mask(self, model_file_type): # Creates segmenter. if model_file_type is ModelFileType.FILE_NAME: base_options = _BaseOptions(model_asset_path=self.model_path) @@ -148,7 +148,7 @@ class ImageSegmenterTest(parameterized.TestCase): # a context. segmenter.close() - def test_succeeds_with_confidence_mask(self): + def test_segment_succeeds_with_confidence_mask(self): # Creates segmenter. base_options = _BaseOptions(model_asset_path=self.model_path) @@ -192,8 +192,7 @@ class ImageSegmenterTest(parameterized.TestCase): segmenter.close() @parameterized.parameters( - (ModelFileType.FILE_NAME,), - (ModelFileType.FILE_CONTENT,)) + (ModelFileType.FILE_NAME), (ModelFileType.FILE_CONTENT)) def test_segment_in_context(self, model_file_type): if model_file_type is ModelFileType.FILE_NAME: base_options = _BaseOptions(model_asset_path=self.model_path) From 91b60da1dcc3a3df43ebcf00ebfe93c12b9f5337 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Oct 2022 05:35:05 -0700 Subject: [PATCH 12/15] Updated name for a test case --- mediapipe/tasks/python/test/vision/image_segmenter_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 054c3e96..b53f301b 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -273,7 +273,7 @@ class ImageSegmenterTest(parameterized.TestCase): r'not initialized with the live stream mode'): segmenter.segment_async(self.test_image, 0) - def test_detect_for_video_with_out_of_order_timestamp(self): + def test_segment_for_video_with_out_of_order_timestamp(self): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), running_mode=_RUNNING_MODE.VIDEO) From 5231a0ad9f7588512119547e271e06a7a54ef7c9 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Fri, 21 Oct 2022 13:34:30 -0700 Subject: [PATCH 13/15] Removed SegmenterOptions dataclasses to enumerate options within ImageSegmenterOptions instead --- mediapipe/tasks/python/components/proto/BUILD | 6 +-- .../components/proto/segmenter_options.py | 50 ------------------- .../test/vision/image_segmenter_test.py | 22 +++----- mediapipe/tasks/python/vision/BUILD | 1 + .../tasks/python/vision/image_segmenter.py | 44 ++++++++++------ 5 files changed, 38 insertions(+), 85 deletions(-) diff --git a/mediapipe/tasks/python/components/proto/BUILD b/mediapipe/tasks/python/components/proto/BUILD index a58f77e6..ef37d927 100644 --- a/mediapipe/tasks/python/components/proto/BUILD +++ b/mediapipe/tasks/python/components/proto/BUILD @@ -20,9 +20,5 @@ licenses(["notice"]) py_library( name = "segmenter_options", - srcs = ["segmenter_options.py"], - deps = [ - "//mediapipe/tasks/cc/components/proto:segmenter_options_py_pb2", - "//mediapipe/tasks/python/core:optional_dependencies", - ], + srcs = ["segmenter_options.py"] ) diff --git a/mediapipe/tasks/python/components/proto/segmenter_options.py b/mediapipe/tasks/python/components/proto/segmenter_options.py index dcf34cc3..5f8e2277 100644 --- a/mediapipe/tasks/python/components/proto/segmenter_options.py +++ b/mediapipe/tasks/python/components/proto/segmenter_options.py @@ -13,14 +13,7 @@ # limitations under the License. """Segmenter options data class.""" -import dataclasses import enum -from typing import Any, Optional - -from mediapipe.tasks.cc.components.proto import segmenter_options_pb2 -from mediapipe.tasks.python.core.optional_dependencies import doc_controls - -_SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions class OutputType(enum.Enum): @@ -33,46 +26,3 @@ class Activation(enum.Enum): NONE = 0 SIGMOID = 1 SOFTMAX = 2 - - -@dataclasses.dataclass -class SegmenterOptions: - """Options for segmentation processor. - Attributes: - output_type: The output mask type allows specifying the type of - post-processing to perform on the raw model results. - activation: Activation function to apply to input tensor. - """ - - output_type: Optional[OutputType] = OutputType.CATEGORY_MASK - activation: Optional[Activation] = Activation.NONE - - @doc_controls.do_not_generate_docs - def to_pb2(self) -> _SegmenterOptionsProto: - """Generates a protobuf object to pass to the C++ layer.""" - return _SegmenterOptionsProto( - output_type=self.output_type.value, - activation=self.activation.value - ) - - @classmethod - @doc_controls.do_not_generate_docs - def create_from_pb2( - cls, pb2_obj: _SegmenterOptionsProto) -> "SegmenterOptions": - """Creates a `SegmenterOptions` object from the given protobuf object.""" - return SegmenterOptions( - output_type=OutputType(pb2_obj.output_type), - activation=Activation(pb2_obj.output_type) - ) - - def __eq__(self, other: Any) -> bool: - """Checks if this object is equal to the given object. - Args: - other: The object to be compared with. - Returns: - True if the objects are equal. - """ - if not isinstance(other, SegmenterOptions): - return False - - return self.to_pb2().__eq__(other.to_pb2()) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index b53f301b..2395eae5 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -35,7 +35,6 @@ _Image = image_module.Image _ImageFormat = image_frame_module.ImageFormat _OutputType = segmenter_options.OutputType _Activation = segmenter_options.Activation -_SegmenterOptions = segmenter_options.SegmenterOptions _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions _RUNNING_MODE = running_mode_module.VisionTaskRunningMode @@ -125,9 +124,8 @@ class ImageSegmenterTest(parameterized.TestCase): # Should never happen raise ValueError('model_file_type is invalid.') - segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions(base_options=base_options, - segmenter_options=segmenter_options) + output_type=_OutputType.CATEGORY_MASK) segmenter = _ImageSegmenter.create_from_options(options) # Performs image segmentation on the input. @@ -153,19 +151,16 @@ class ImageSegmenterTest(parameterized.TestCase): base_options = _BaseOptions(model_asset_path=self.model_path) # Run segmentation on the model in CATEGORY_MASK mode. - segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions(base_options=base_options, - segmenter_options=segmenter_options) + output_type=_OutputType.CATEGORY_MASK) segmenter = _ImageSegmenter.create_from_options(options) category_masks = segmenter.segment(self.test_image) category_mask = category_masks[0].numpy_view() # Run segmentation on the model in CONFIDENCE_MASK mode. - segmenter_options = _SegmenterOptions( - output_type=_OutputType.CONFIDENCE_MASK, - activation=_Activation.SOFTMAX) options = _ImageSegmenterOptions(base_options=base_options, - segmenter_options=segmenter_options) + output_type=_OutputType.CONFIDENCE_MASK, + activation=_Activation.SOFTMAX) segmenter = _ImageSegmenter.create_from_options(options) confidence_masks = segmenter.segment(self.test_image) @@ -204,9 +199,8 @@ class ImageSegmenterTest(parameterized.TestCase): # Should never happen raise ValueError('model_file_type is invalid.') - segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions(base_options=base_options, - segmenter_options=segmenter_options) + output_type=_OutputType.CATEGORY_MASK) with _ImageSegmenter.create_from_options(options) as segmenter: # Performs image segmentation on the input. category_masks = segmenter.segment(self.test_image) @@ -284,10 +278,9 @@ class ImageSegmenterTest(parameterized.TestCase): segmenter.segment_for_video(self.test_image, 0) def test_segment_for_video(self): - segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), - segmenter_options=segmenter_options, + output_type=_OutputType.CATEGORY_MASK, running_mode=_RUNNING_MODE.VIDEO) with _ImageSegmenter.create_from_options(options) as segmenter: for timestamp in range(0, 300, 30): @@ -348,10 +341,9 @@ class ImageSegmenterTest(parameterized.TestCase): self.assertLess(observed_timestamp_ms, timestamp_ms) self.observed_timestamp_ms = timestamp_ms - segmenter_options = _SegmenterOptions(output_type=_OutputType.CATEGORY_MASK) options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), - segmenter_options=segmenter_options, + output_type=_OutputType.CATEGORY_MASK, running_mode=_RUNNING_MODE.LIVE_STREAM, result_callback=check_result) with _ImageSegmenter.create_from_options(options) as segmenter: diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index 3875ea5d..863312e4 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -46,6 +46,7 @@ py_library( "//mediapipe/python:_framework_bindings", "//mediapipe/python:packet_creator", "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/components/proto:segmenter_options_py_pb2", "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_py_pb2", "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/core:base_options", diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 51f80292..e7278eb9 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -19,22 +19,25 @@ from typing import Callable, List, Mapping, Optional from mediapipe.python import packet_creator from mediapipe.python import packet_getter from mediapipe.python._framework_bindings import image as image_module -from mediapipe.python._framework_bindings import packet as packet_module -from mediapipe.python._framework_bindings import task_runner as task_runner_module +from mediapipe.python._framework_bindings import packet +from mediapipe.python._framework_bindings import task_runner +from mediapipe.tasks.cc.components.proto import segmenter_options_pb2 from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_options_pb2 from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.core import task_info as task_info_module from mediapipe.tasks.python.core.optional_dependencies import doc_controls from mediapipe.tasks.python.vision.core import base_vision_task_api -from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module +from mediapipe.tasks.python.vision.core import vision_task_running_mode _BaseOptions = base_options_module.BaseOptions +_SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions _ImageSegmenterOptionsProto = image_segmenter_options_pb2.ImageSegmenterOptions -_SegmenterOptions = segmenter_options.SegmenterOptions -_RunningMode = running_mode_module.VisionTaskRunningMode +_OutputType = segmenter_options.OutputType +_Activation = segmenter_options.Activation +_RunningMode = vision_task_running_mode.VisionTaskRunningMode _TaskInfo = task_info_module.TaskInfo -_TaskRunner = task_runner_module.TaskRunner +_TaskRunner = task_runner.TaskRunner _SEGMENTATION_OUT_STREAM_NAME = 'segmented_mask_out' _SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' @@ -57,14 +60,17 @@ class ImageSegmenterOptions: 2) The video mode for segmenting objects on the decoded frames of a video. 3) The live stream mode for segmenting objects on a live stream of input data, such as from camera. - segmenter_options: Options for the image segmenter task. + output_type: The output mask type allows specifying the type of + post-processing to perform on the raw model results. + activation: Activation function to apply to input tensor. result_callback: The user-defined result callback for processing live stream data. The result callback should only be specified when the running mode is set to the live stream mode. """ base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE - segmenter_options: _SegmenterOptions = _SegmenterOptions() + output_type: Optional[_OutputType] = _OutputType.CATEGORY_MASK + activation: Optional[_Activation] = _Activation.NONE result_callback: Optional[ Callable[[List[image_module.Image], image_module.Image, int], None]] = None @@ -74,8 +80,10 @@ class ImageSegmenterOptions: """Generates an ImageSegmenterOptions protobuf object.""" base_options_proto = self.base_options.to_pb2() base_options_proto.use_stream_mode = False if self.running_mode == _RunningMode.IMAGE else True - segmenter_options_proto = self.segmenter_options.to_pb2() - + segmenter_options_proto = _SegmenterOptionsProto( + output_type=self.output_type.value, + activation=self.activation.value + ) return _ImageSegmenterOptionsProto( base_options=base_options_proto, segmenter_options=segmenter_options_proto @@ -127,7 +135,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): RuntimeError: If other types of error occurred. """ - def packets_callback(output_packets: Mapping[str, packet_module.Packet]): + def packets_callback(output_packets: Mapping[str, packet.Packet]): if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty(): return segmentation_result = packet_getter.get_image_list( @@ -159,8 +167,11 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): image: MediaPipe Image. Returns: - A segmentation result object that contains a list of segmentation masks - as images. + If the output_type is CATEGORY_MASK, the returned vector of images is + per-category segmented image mask. + If the output_type is CONFIDENCE_MASK, the returned vector of images + contains only one confidence image mask. A segmentation result object that + contains a list of segmentation masks as images. Raises: ValueError: If any of the input arguments is invalid. @@ -186,8 +197,11 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): timestamp_ms: The timestamp of the input video frame in milliseconds. Returns: - A segmentation result object that contains a list of segmentation masks - as images. + If the output_type is CATEGORY_MASK, the returned vector of images is + per-category segmented image mask. + If the output_type is CONFIDENCE_MASK, the returned vector of images + contains only one confidence image mask. A segmentation result object that + contains a list of segmentation masks as images. Raises: ValueError: If any of the input arguments is invalid. From 024a6866a7ba41bb0e8e33106c2195f0ced88e6e Mon Sep 17 00:00:00 2001 From: kinaryml Date: Fri, 21 Oct 2022 13:39:59 -0700 Subject: [PATCH 14/15] Removed some unneeded aliases --- .../tasks/python/test/vision/image_segmenter_test.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 2395eae5..9b4e2c4b 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -23,21 +23,21 @@ from absl.testing import absltest from absl.testing import parameterized from mediapipe.python._framework_bindings import image as image_module -from mediapipe.python._framework_bindings import image_frame as image_frame_module +from mediapipe.python._framework_bindings import image_frame from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.test import test_utils from mediapipe.tasks.python.vision import image_segmenter -from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module +from mediapipe.tasks.python.vision.core import vision_task_running_mode _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image -_ImageFormat = image_frame_module.ImageFormat +_ImageFormat = image_frame.ImageFormat _OutputType = segmenter_options.OutputType _Activation = segmenter_options.Activation _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions -_RUNNING_MODE = running_mode_module.VisionTaskRunningMode +_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode _MODEL_FILE = 'deeplabv3.tflite' _IMAGE_FILE = 'segmentation_input_rotation0.jpg' From ebb2686fb4d259ea61251ca23a4c4dc4705344e1 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Sat, 22 Oct 2022 03:34:26 -0700 Subject: [PATCH 15/15] Moved OutputType and Activation to image_segmenter --- mediapipe/tasks/python/components/proto/BUILD | 24 ---------------- .../components/proto/segmenter_options.py | 28 ------------------- mediapipe/tasks/python/test/vision/BUILD | 1 - .../test/vision/image_segmenter_test.py | 5 ++-- mediapipe/tasks/python/vision/BUILD | 1 - .../tasks/python/vision/image_segmenter.py | 20 +++++++++---- 6 files changed, 17 insertions(+), 62 deletions(-) delete mode 100644 mediapipe/tasks/python/components/proto/BUILD delete mode 100644 mediapipe/tasks/python/components/proto/segmenter_options.py diff --git a/mediapipe/tasks/python/components/proto/BUILD b/mediapipe/tasks/python/components/proto/BUILD deleted file mode 100644 index ef37d927..00000000 --- a/mediapipe/tasks/python/components/proto/BUILD +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright 2022 The MediaPipe Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Placeholder for internal Python strict library compatibility macro. - -package(default_visibility = ["//mediapipe/tasks:internal"]) - -licenses(["notice"]) - -py_library( - name = "segmenter_options", - srcs = ["segmenter_options.py"] -) diff --git a/mediapipe/tasks/python/components/proto/segmenter_options.py b/mediapipe/tasks/python/components/proto/segmenter_options.py deleted file mode 100644 index 5f8e2277..00000000 --- a/mediapipe/tasks/python/components/proto/segmenter_options.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright 2022 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Segmenter options data class.""" - -import enum - - -class OutputType(enum.Enum): - UNSPECIFIED = 0 - CATEGORY_MASK = 1 - CONFIDENCE_MASK = 2 - - -class Activation(enum.Enum): - NONE = 0 - SIGMOID = 1 - SOFTMAX = 2 diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 63fc56b4..321b33a6 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -48,7 +48,6 @@ py_test( "//mediapipe/python:_framework_bindings", "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/test:test_utils", - "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/vision:image_segmenter", "//mediapipe/tasks/python/vision/core:vision_task_running_mode", "@absl_py//absl/testing:parameterized", diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 9b4e2c4b..3fa01e62 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -24,7 +24,6 @@ from absl.testing import parameterized from mediapipe.python._framework_bindings import image as image_module from mediapipe.python._framework_bindings import image_frame -from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.test import test_utils from mediapipe.tasks.python.vision import image_segmenter @@ -33,8 +32,8 @@ from mediapipe.tasks.python.vision.core import vision_task_running_mode _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _ImageFormat = image_frame.ImageFormat -_OutputType = segmenter_options.OutputType -_Activation = segmenter_options.Activation +_OutputType = image_segmenter.OutputType +_Activation = image_segmenter.Activation _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions _RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index 863312e4..da9072f1 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -48,7 +48,6 @@ py_library( "//mediapipe/python:packet_getter", "//mediapipe/tasks/cc/components/proto:segmenter_options_py_pb2", "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_options_py_pb2", - "//mediapipe/tasks/python/components/proto:segmenter_options", "//mediapipe/tasks/python/core:base_options", "//mediapipe/tasks/python/core:optional_dependencies", "//mediapipe/tasks/python/core:task_info", diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index e7278eb9..aea5a855 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -15,6 +15,7 @@ import dataclasses from typing import Callable, List, Mapping, Optional +import enum from mediapipe.python import packet_creator from mediapipe.python import packet_getter @@ -23,7 +24,6 @@ from mediapipe.python._framework_bindings import packet from mediapipe.python._framework_bindings import task_runner from mediapipe.tasks.cc.components.proto import segmenter_options_pb2 from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_options_pb2 -from mediapipe.tasks.python.components.proto import segmenter_options from mediapipe.tasks.python.core import base_options as base_options_module from mediapipe.tasks.python.core import task_info as task_info_module from mediapipe.tasks.python.core.optional_dependencies import doc_controls @@ -33,8 +33,6 @@ from mediapipe.tasks.python.vision.core import vision_task_running_mode _BaseOptions = base_options_module.BaseOptions _SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions _ImageSegmenterOptionsProto = image_segmenter_options_pb2.ImageSegmenterOptions -_OutputType = segmenter_options.OutputType -_Activation = segmenter_options.Activation _RunningMode = vision_task_running_mode.VisionTaskRunningMode _TaskInfo = task_info_module.TaskInfo _TaskRunner = task_runner.TaskRunner @@ -48,6 +46,18 @@ _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph' _MICRO_SECONDS_PER_MILLISECOND = 1000 +class OutputType(enum.Enum): + UNSPECIFIED = 0 + CATEGORY_MASK = 1 + CONFIDENCE_MASK = 2 + + +class Activation(enum.Enum): + NONE = 0 + SIGMOID = 1 + SOFTMAX = 2 + + @dataclasses.dataclass class ImageSegmenterOptions: """Options for the image segmenter task. @@ -69,8 +79,8 @@ class ImageSegmenterOptions: """ base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE - output_type: Optional[_OutputType] = _OutputType.CATEGORY_MASK - activation: Optional[_Activation] = _Activation.NONE + output_type: Optional[OutputType] = OutputType.CATEGORY_MASK + activation: Optional[Activation] = Activation.NONE result_callback: Optional[ Callable[[List[image_module.Image], image_module.Image, int], None]] = None