From a03fa448dc09d75da7f5eecae16a1a2909a46df6 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 13 Apr 2023 11:55:37 -0700 Subject: [PATCH] Explicitly state the modes in the tests for ImageSegmenterOptions and InteractiveSegmenterOptions --- .../test/vision/image_segmenter_test.py | 32 +++++++++++++------ .../test/vision/interactive_segmenter_test.py | 18 ++++++++--- 2 files changed, 36 insertions(+), 14 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 327925191..aa557281f 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -157,7 +157,9 @@ class ImageSegmenterTest(parameterized.TestCase): raise ValueError('model_file_type is invalid.') options = _ImageSegmenterOptions( - base_options=base_options, output_category_mask=True) + base_options=base_options, output_category_mask=True, + output_confidence_masks=False + ) segmenter = _ImageSegmenter.create_from_options(options) # Performs image segmentation on the input. @@ -188,8 +190,9 @@ class ImageSegmenterTest(parameterized.TestCase): # Run segmentation on the model in CONFIDENCE_MASK mode. options = _ImageSegmenterOptions( - base_options=base_options, - activation=_Activation.SOFTMAX) + base_options=base_options, output_category_mask=False, + output_confidence_masks=True, activation=_Activation.SOFTMAX + ) with _ImageSegmenter.create_from_options(options) as segmenter: segmentation_result = segmenter.segment(test_image) @@ -279,7 +282,9 @@ class ImageSegmenterTest(parameterized.TestCase): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), output_category_mask=True, - running_mode=_RUNNING_MODE.VIDEO) + output_confidence_masks=False, + running_mode=_RUNNING_MODE.VIDEO + ) with _ImageSegmenter.create_from_options(options) as segmenter: for timestamp in range(0, 300, 30): segmentation_result = segmenter.segment_for_video( @@ -297,8 +302,10 @@ class ImageSegmenterTest(parameterized.TestCase): os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))) options = _ImageSegmenterOptions( - base_options=_BaseOptions(model_asset_path=self.model_path), - running_mode=_RUNNING_MODE.VIDEO) + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.VIDEO, output_category_mask=False, + output_confidence_masks=True + ) with _ImageSegmenter.create_from_options(options) as segmenter: for timestamp in range(0, 300, 30): segmentation_result = segmenter.segment_for_video( @@ -370,8 +377,10 @@ class ImageSegmenterTest(parameterized.TestCase): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), output_category_mask=True, + output_confidence_masks=False, running_mode=_RUNNING_MODE.LIVE_STREAM, - result_callback=check_result) + 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) @@ -405,9 +414,12 @@ class ImageSegmenterTest(parameterized.TestCase): self.observed_timestamp_ms = timestamp_ms options = _ImageSegmenterOptions( - base_options=_BaseOptions(model_asset_path=self.model_path), - running_mode=_RUNNING_MODE.LIVE_STREAM, - result_callback=check_result) + base_options=_BaseOptions(model_asset_path=self.model_path), + running_mode=_RUNNING_MODE.LIVE_STREAM, + output_category_mask=False, + output_confidence_masks=True, + result_callback=check_result + ) with _ImageSegmenter.create_from_options(options) as segmenter: for timestamp in range(0, 300, 30): segmenter.segment_async(test_image, timestamp) diff --git a/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py b/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py index 6af15aa09..aea4f8a1d 100644 --- a/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py @@ -200,7 +200,8 @@ class InteractiveSegmenterTest(parameterized.TestCase): raise ValueError('model_file_type is invalid.') options = _InteractiveSegmenterOptions( - base_options=base_options, output_category_mask=True + base_options=base_options, output_category_mask=True, + output_confidence_masks=False ) segmenter = _InteractiveSegmenter.create_from_options(options) @@ -252,7 +253,10 @@ class InteractiveSegmenterTest(parameterized.TestCase): roi = _RegionOfInterest(format=roi_format, keypoint=keypoint) # Run segmentation on the model in CONFIDENCE_MASK mode. - options = _InteractiveSegmenterOptions(base_options=base_options) + options = _InteractiveSegmenterOptions( + base_options=base_options, output_category_mask=False, + output_confidence_masks=True + ) with _InteractiveSegmenter.create_from_options(options) as segmenter: # Perform segmentation @@ -284,7 +288,10 @@ class InteractiveSegmenterTest(parameterized.TestCase): ) # Run segmentation on the model in CONFIDENCE_MASK mode. - options = _InteractiveSegmenterOptions(base_options=base_options) + options = _InteractiveSegmenterOptions( + base_options=base_options, output_category_mask=False, + output_confidence_masks=True + ) with _InteractiveSegmenter.create_from_options(options) as segmenter: # Perform segmentation @@ -310,7 +317,10 @@ class InteractiveSegmenterTest(parameterized.TestCase): ) # Run segmentation on the model in CONFIDENCE_MASK mode. - options = _InteractiveSegmenterOptions(base_options=base_options) + options = _InteractiveSegmenterOptions( + base_options=base_options, output_category_mask=False, + output_confidence_masks=True + ) with self.assertRaisesRegex( ValueError, "This task doesn't support region-of-interest."