From 723cb2a91977818c2198cf394bd81db8bba1fc25 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Apr 2023 02:49:13 -0700 Subject: [PATCH 1/5] Populate labels using model metadata for the ImageSegmenter Python API --- .../test/vision/image_segmenter_test.py | 33 ++++++++++ mediapipe/tasks/python/vision/BUILD | 1 + .../tasks/python/vision/image_segmenter.py | 63 +++++++++++++++++++ 3 files changed, 97 insertions(+) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 7f0b47eb..d993315e 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -45,6 +45,29 @@ _SEGMENTATION_FILE = 'segmentation_golden_rotation0.png' _MASK_MAGNIFICATION_FACTOR = 10 _MASK_SIMILARITY_THRESHOLD = 0.98 _TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision' +_EXPECTED_LABELS = [ + "background", + "aeroplane", + "bicycle", + "bird", + "boat", + "bottle", + "bus", + "car", + "cat", + "chair", + "cow", + "dining table", + "dog", + "horse", + "motorbike", + "person", + "potted plant", + "sheep", + "sofa", + "train", + "tv" +] def _similar_to_uint8_mask(actual_mask, expected_mask): @@ -214,6 +237,16 @@ class ImageSegmenterTest(parameterized.TestCase): f'Number of pixels in the candidate mask differing from that of the ' f'ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.') + def test_get_labels_succeeds(self): + expected_labels = _EXPECTED_LABELS + base_options = _BaseOptions(model_asset_path=self.model_path) + options = _ImageSegmenterOptions( + base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + with _ImageSegmenter.create_from_options(options) as segmenter: + # Performs image segmentation on the input. + actual_labels = segmenter.get_labels() + self.assertListEqual(actual_labels, expected_labels) + def test_missing_result_callback(self): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), diff --git a/mediapipe/tasks/python/vision/BUILD b/mediapipe/tasks/python/vision/BUILD index 046ce2dc..71675779 100644 --- a/mediapipe/tasks/python/vision/BUILD +++ b/mediapipe/tasks/python/vision/BUILD @@ -71,6 +71,7 @@ py_library( "//mediapipe/python:_framework_bindings", "//mediapipe/python:packet_creator", "//mediapipe/python:packet_getter", + "//mediapipe/tasks/cc/vision/image_segmenter/calculators:tensors_to_segmentation_calculator_py_pb2", "//mediapipe/tasks/cc/vision/image_segmenter/proto:image_segmenter_graph_options_py_pb2", "//mediapipe/tasks/cc/vision/image_segmenter/proto:segmenter_options_py_pb2", "//mediapipe/tasks/python/components/containers:rect", diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index e50ffbf7..f70f1653 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -21,6 +21,7 @@ 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 +from mediapipe.tasks.cc.vision.image_segmenter.calculators import tensors_to_segmentation_calculator_pb2 from mediapipe.tasks.cc.vision.image_segmenter.proto import image_segmenter_graph_options_pb2 from mediapipe.tasks.cc.vision.image_segmenter.proto import segmenter_options_pb2 from mediapipe.tasks.python.components.containers import rect @@ -38,6 +39,9 @@ _SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions _ImageSegmenterGraphOptionsProto = ( image_segmenter_graph_options_pb2.ImageSegmenterGraphOptions ) +TensorsToSegmentationCalculatorOptionsProto = ( + tensors_to_segmentation_calculator_pb2.TensorsToSegmentationCalculatorOptions +) _RunningMode = vision_task_running_mode.VisionTaskRunningMode _ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions _TaskInfo = task_info_module.TaskInfo @@ -49,6 +53,7 @@ _IMAGE_OUT_STREAM_NAME = 'image_out' _IMAGE_TAG = 'IMAGE' _NORM_RECT_STREAM_NAME = 'norm_rect_in' _NORM_RECT_TAG = 'NORM_RECT' +_TENSORS_TO_SEGMENTATION_CALCULATOR_NAME = 'mediapipe.tasks.TensorsToSegmentationCalculator' _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph' _MICRO_SECONDS_PER_MILLISECOND = 1000 @@ -130,6 +135,40 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): An example of such model can be found at: https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2 """ + def __init__(self, graph_config, running_mode, packet_callback): + super(ImageSegmenter, self).__init__( + graph_config, running_mode, packet_callback + ) + self._populate_labels() + + def _populate_labels(self): + """ + Populate the labelmap in TensorsToSegmentationCalculator to labels field. + + Returns: + Exception if there is an error during finding TensorsToSegmentationCalculator. + :return: + """ + self.labels = [] + graph_config = self._runner.get_graph_config() + found_tensors_to_segmentation = False + + for node in graph_config.node: + if _TENSORS_TO_SEGMENTATION_CALCULATOR_NAME in node.name: + if found_tensors_to_segmentation: + raise Exception( + f"The graph has more than one " + f"{_TENSORS_TO_SEGMENTATION_CALCULATOR_NAME}." + ) + found_tensors_to_segmentation = True + options = node.options.Extensions[ + TensorsToSegmentationCalculatorOptionsProto.ext + ] + if options.label_items: + for i in range(len(options.label_items)): + if i not in options.label_items: + raise Exception(f"The labelmap has no expected key: {i}.") + self.labels.append(options.label_items[i].name) @classmethod def create_from_model_path(cls, model_path: str) -> 'ImageSegmenter': @@ -209,6 +248,30 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): packets_callback if options.result_callback else None, ) + def get_labels(self): + """ Get the category label list of the ImageSegmenter can recognize. + + For CATEGORY_MASK type, the index in the category mask corresponds to the + category in the label list. + For CONFIDENCE_MASK type, the output mask list at index corresponds to the + category in the label list. + + If there is no label map provided in the model file, empty label list is + returned. + + Returns: + 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. + RuntimeError: If image segmentation failed to run. + """ + return self.labels + def segment( self, image: image_module.Image, From 1919b0e34125a3e9a4287550b9d9a357fac50f96 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Apr 2023 02:54:03 -0700 Subject: [PATCH 2/5] Updated docstrings for get_labels --- mediapipe/tasks/python/vision/image_segmenter.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index f70f1653..4f57b89d 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -258,17 +258,6 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): If there is no label map provided in the model file, empty label list is returned. - - Returns: - 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. - RuntimeError: If image segmentation failed to run. """ return self.labels From 1cb404bea16a4f36df8abeb583a72b9819776583 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Apr 2023 21:31:14 -0700 Subject: [PATCH 3/5] Changed labels to be a property --- .../test/vision/image_segmenter_test.py | 8 +++-- .../tasks/python/vision/image_segmenter.py | 31 ++++++++++--------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index b54b5399..3458bb50 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -247,14 +247,16 @@ class ImageSegmenterTest(parameterized.TestCase): ) ) - def test_get_labels_succeeds(self): + def test_labels_succeeds(self): expected_labels = _EXPECTED_LABELS base_options = _BaseOptions(model_asset_path=self.model_path) options = _ImageSegmenterOptions( - base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + base_options=base_options, output_category_mask=True, + output_confidence_masks=False + ) with _ImageSegmenter.create_from_options(options) as segmenter: # Performs image segmentation on the input. - actual_labels = segmenter.get_labels() + actual_labels = segmenter.labels self.assertListEqual(actual_labels, expected_labels) def test_missing_result_callback(self): diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 4119f263..a6c9501c 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -151,7 +151,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): Exception if there is an error during finding TensorsToSegmentationCalculator. :return: """ - self.labels = [] + self._labels = [] graph_config = self._runner.get_graph_config() found_tensors_to_segmentation = False @@ -170,7 +170,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): for i in range(len(options.label_items)): if i not in options.label_items: raise Exception(f"The labelmap has no expected key: {i}.") - self.labels.append(options.label_items[i].name) + self._labels.append(options.label_items[i].name) @classmethod def create_from_model_path(cls, model_path: str) -> 'ImageSegmenter': @@ -271,19 +271,6 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): packets_callback if options.result_callback else None, ) - def get_labels(self): - """ Get the category label list of the ImageSegmenter can recognize. - - For CATEGORY_MASK type, the index in the category mask corresponds to the - category in the label list. - For CONFIDENCE_MASK type, the output mask list at index corresponds to the - category in the label list. - - If there is no label map provided in the model file, empty label list is - returned. - """ - return self.labels - def segment( self, image: image_module.Image, @@ -427,3 +414,17 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): normalized_rect.to_pb2() ).at(timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), }) + + @property + def labels(self) -> List[str]: + """ Get the category label list of the ImageSegmenter can recognize. + + For CATEGORY_MASK type, the index in the category mask corresponds to the + category in the label list. + For CONFIDENCE_MASK type, the output mask list at index corresponds to the + category in the label list. + + If there is no label map provided in the model file, empty label list is + returned. + """ + return self._labels From 67b72e4fe9b6765c3d134d88a6ba77ac50a35a05 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Apr 2023 21:43:38 -0700 Subject: [PATCH 4/5] Code cleanup --- .../python/test/vision/image_segmenter_test.py | 7 ++++--- mediapipe/tasks/python/vision/image_segmenter.py | 15 ++++++++------- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 3458bb50..009dc685 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -247,12 +247,13 @@ class ImageSegmenterTest(parameterized.TestCase): ) ) - def test_labels_succeeds(self): + @parameterized.parameters((True, False), (False, True)) + def test_labels_succeeds(self, output_category_mask, output_confidence_masks): expected_labels = _EXPECTED_LABELS base_options = _BaseOptions(model_asset_path=self.model_path) options = _ImageSegmenterOptions( - base_options=base_options, output_category_mask=True, - output_confidence_masks=False + base_options=base_options, output_category_mask=output_category_mask, + output_confidence_masks=output_confidence_masks ) with _ImageSegmenter.create_from_options(options) as segmenter: # Performs image segmentation on the input. diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index a6c9501c..220d7818 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -129,27 +129,28 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): Output tensors: (kTfLiteUInt8/kTfLiteFloat32) - list of segmented masks. - - if `output_type` is CATEGORY_MASK, uint8 Image, Image vector of size 1. - - if `output_type` is CONFIDENCE_MASK, float32 Image list of size + - if `output_category_mask` is True, uint8 Image, Image vector of size 1. + - if `output_confidence_masks` is True, float32 Image list of size `channels`. - batch is always 1 An example of such model can be found at: https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2 """ - def __init__(self, graph_config, running_mode, packet_callback): + def __init__(self, graph_config, running_mode, packet_callback) -> None: + """Initializes the `ImageSegmenter` object.""" super(ImageSegmenter, self).__init__( graph_config, running_mode, packet_callback ) self._populate_labels() - def _populate_labels(self): + def _populate_labels(self) -> None: """ Populate the labelmap in TensorsToSegmentationCalculator to labels field. - Returns: - Exception if there is an error during finding TensorsToSegmentationCalculator. - :return: + Raises: + Exception if there is an error during finding + TensorsToSegmentationCalculator. """ self._labels = [] graph_config = self._runner.get_graph_config() From a1aab66c8d0635c2f9d26f6dc41c7517c6e82dfc Mon Sep 17 00:00:00 2001 From: kinaryml Date: Tue, 18 Apr 2023 21:50:27 -0700 Subject: [PATCH 5/5] Fixed a typo in docstring --- mediapipe/tasks/python/vision/image_segmenter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 220d7818..077f7d28 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -418,7 +418,7 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): @property def labels(self) -> List[str]: - """ Get the category label list of the ImageSegmenter can recognize. + """ Get the category label list the ImageSegmenter can recognize. For CATEGORY_MASK type, the index in the category mask corresponds to the category in the label list.