From 3f68f90238a36521994c9ab5ee8500473f7f657d Mon Sep 17 00:00:00 2001 From: kinaryml Date: Wed, 12 Apr 2023 14:37:16 -0700 Subject: [PATCH 1/9] Deprecated output_type for the ImageSegmenter and InteractiveSegmenter APIs --- mediapipe/tasks/python/test/vision/BUILD | 19 ++ .../test/vision/image_segmenter_test.py | 194 ++++++++++++------ .../test/vision/interactive_segmenter_test.py | 29 ++- .../tasks/python/vision/image_segmenter.py | 96 ++++++--- .../python/vision/interactive_segmenter.py | 64 ++++-- mediapipe/tasks/testdata/vision/BUILD | 2 + 6 files changed, 272 insertions(+), 132 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/BUILD b/mediapipe/tasks/python/test/vision/BUILD index 704e1af5..46838c26 100644 --- a/mediapipe/tasks/python/test/vision/BUILD +++ b/mediapipe/tasks/python/test/vision/BUILD @@ -93,6 +93,25 @@ py_test( ], ) +py_test( + name = "interactive_segmenter_test", + srcs = ["interactive_segmenter_test.py"], + data = [ + "//mediapipe/tasks/testdata/vision:test_images", + "//mediapipe/tasks/testdata/vision:test_models", + ], + deps = [ + "//mediapipe/python:_framework_bindings", + "//mediapipe/tasks/python/components/containers:keypoint", + "//mediapipe/tasks/python/components/containers:rect", + "//mediapipe/tasks/python/core:base_options", + "//mediapipe/tasks/python/test:test_utils", + "//mediapipe/tasks/python/vision:interactive_segmenter", + "//mediapipe/tasks/python/vision/core:image_processing_options", + "//mediapipe/tasks/python/vision/core:vision_task_running_mode", + ], +) + py_test( name = "face_detector_test", srcs = ["face_detector_test.py"], diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index 7f0b47eb..32792519 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -30,10 +30,10 @@ 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 +ImageSegmenterResult = image_segmenter.ImageSegmenterResult _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _ImageFormat = image_frame.ImageFormat -_OutputType = image_segmenter.ImageSegmenterOptions.OutputType _Activation = image_segmenter.ImageSegmenterOptions.Activation _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions @@ -42,11 +42,33 @@ _RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode _MODEL_FILE = 'deeplabv3.tflite' _IMAGE_FILE = 'segmentation_input_rotation0.jpg' _SEGMENTATION_FILE = 'segmentation_golden_rotation0.png' +_CAT_IMAGE = 'cat.jpg' +_CAT_MASK = 'cat_mask.jpg' _MASK_MAGNIFICATION_FACTOR = 10 _MASK_SIMILARITY_THRESHOLD = 0.98 _TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision' +def _calculate_soft_iou(m1, m2): + intersection_sum = np.sum(m1 * m2) + union_sum = np.sum(m1 * m1) + np.sum(m2 * m2) - intersection_sum + + if union_sum > 0: + return intersection_sum / union_sum + else: + return 0 + + +def _similar_to_float_mask(actual_mask, expected_mask, similarity_threshold): + actual_mask = actual_mask.numpy_view() + expected_mask = expected_mask.numpy_view() / 255.0 + + return ( + actual_mask.shape == expected_mask.shape + and _calculate_soft_iou(actual_mask, expected_mask) > similarity_threshold + ) + + 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() @@ -84,6 +106,14 @@ class ImageSegmenterTest(parameterized.TestCase): self.model_path = test_utils.get_test_data_path( os.path.join(_TEST_DATA_DIR, _MODEL_FILE)) + def _load_segmentation_mask(self, file_path: str): + # Loads ground truth segmentation file. + gt_segmentation_data = cv2.imread( + test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, file_path)), + cv2.IMREAD_GRAYSCALE, + ) + return _Image(_ImageFormat.GRAY8, gt_segmentation_data) + 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: @@ -127,20 +157,19 @@ class ImageSegmenterTest(parameterized.TestCase): raise ValueError('model_file_type is invalid.') options = _ImageSegmenterOptions( - base_options=base_options, output_type=_OutputType.CATEGORY_MASK) + base_options=base_options, output_category_mask=True) segmenter = _ImageSegmenter.create_from_options(options) # Performs image segmentation on the input. - category_masks = segmenter.segment(self.test_image) - self.assertLen(category_masks, 1) - category_mask = category_masks[0] + segmentation_result = segmenter.segment(self.test_image) + category_mask = segmentation_result.category_mask result_pixels = category_mask.numpy_view().flatten() # Check if data type of `category_mask` is correct. self.assertEqual(result_pixels.dtype, np.uint8) self.assertTrue( - _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + _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}.') @@ -152,67 +181,33 @@ class ImageSegmenterTest(parameterized.TestCase): # Creates segmenter. base_options = _BaseOptions(model_asset_path=self.model_path) - # Run segmentation on the model in CATEGORY_MASK mode. - options = _ImageSegmenterOptions( - base_options=base_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() + # Load the cat image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))) # Run segmentation on the model in CONFIDENCE_MASK mode. options = _ImageSegmenterOptions( base_options=base_options, - output_type=_OutputType.CONFIDENCE_MASK, activation=_Activation.SOFTMAX) - segmenter = _ImageSegmenter.create_from_options(options) - confidence_masks = segmenter.segment(self.test_image) - # Check if confidence mask shape is correct. - self.assertLen( - 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. - 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.') - - 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. - category_masks = segmenter.segment(self.test_image) - self.assertLen(category_masks, 1) + segmentation_result = segmenter.segment(test_image) + confidence_masks = segmentation_result.confidence_masks + + # Check if confidence mask shape is correct. + self.assertLen( + confidence_masks, 21, + 'Number of confidence masks must match with number of categories.') + + # Loads ground truth segmentation file. + expected_mask = self._load_segmentation_mask(_CAT_MASK) 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}.') + _similar_to_float_mask( + confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD + ) + ) def test_missing_result_callback(self): options = _ImageSegmenterOptions( @@ -280,20 +275,49 @@ class ImageSegmenterTest(parameterized.TestCase): ValueError, r'Input timestamp must be monotonically increasing'): segmenter.segment_for_video(self.test_image, 0) - def test_segment_for_video(self): + def test_segment_for_video_in_category_mask_mode(self): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), - output_type=_OutputType.CATEGORY_MASK, + output_category_mask=True, 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.assertLen(category_masks, 1) + segmentation_result = segmenter.segment_for_video( + self.test_image, timestamp) + category_mask = segmentation_result.category_mask self.assertTrue( - _similar_to_uint8_mask(category_masks[0], self.test_seg_image), + _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}.') + def test_segment_for_video_in_confidence_mask_mode(self): + # Load the cat image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))) + + options = _ImageSegmenterOptions( + base_options=_BaseOptions(model_asset_path=self.model_path), + 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( + test_image, timestamp) + confidence_masks = segmentation_result.confidence_masks + + # Check if confidence mask shape is correct. + self.assertLen( + confidence_masks, 21, + 'Number of confidence masks must match with number of categories.') + + # Loads ground truth segmentation file. + expected_mask = self._load_segmentation_mask(_CAT_MASK) + self.assertTrue( + _similar_to_float_mask( + confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD + ) + ) + def test_calling_segment_in_live_stream_mode(self): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), @@ -325,13 +349,13 @@ class ImageSegmenterTest(parameterized.TestCase): ValueError, r'Input timestamp must be monotonically increasing'): segmenter.segment_async(self.test_image, 0) - def test_segment_async_calls(self): + def test_segment_async_calls_in_category_mask_mode(self): observed_timestamp_ms = -1 - def check_result(result: List[image_module.Image], output_image: _Image, + def check_result(result: ImageSegmenterResult, output_image: _Image, timestamp_ms: int): # Get the output category mask. - category_mask = result[0] + category_mask = result.category_mask 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) @@ -345,13 +369,49 @@ class ImageSegmenterTest(parameterized.TestCase): options = _ImageSegmenterOptions( base_options=_BaseOptions(model_asset_path=self.model_path), - output_type=_OutputType.CATEGORY_MASK, + output_category_mask=True, 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) + def test_segment_async_calls_in_confidence_mask_mode(self): + # Load the cat image. + test_image = _Image.create_from_file( + test_utils.get_test_data_path( + os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))) + + # Loads ground truth segmentation file. + expected_mask = self._load_segmentation_mask(_CAT_MASK) + observed_timestamp_ms = -1 + + def check_result(result: ImageSegmenterResult, output_image: _Image, + timestamp_ms: int): + # Get the output category mask. + confidence_masks = result.confidence_masks + + # Check if confidence mask shape is correct. + self.assertLen( + confidence_masks, 21, + 'Number of confidence masks must match with number of categories.') + + self.assertTrue( + _similar_to_float_mask( + confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD + ) + ) + self.assertLess(observed_timestamp_ms, timestamp_ms) + 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) + with _ImageSegmenter.create_from_options(options) as segmenter: + for timestamp in range(0, 300, 30): + segmenter.segment_async(test_image, timestamp) + if __name__ == '__main__': absltest.main() diff --git a/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py b/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py index e8c52ae3..6af15aa0 100644 --- a/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/interactive_segmenter_test.py @@ -30,12 +30,12 @@ from mediapipe.tasks.python.test import test_utils from mediapipe.tasks.python.vision import interactive_segmenter from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module +InteractiveSegmenterResult = interactive_segmenter.InteractiveSegmenterResult _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _ImageFormat = image_frame.ImageFormat _NormalizedKeypoint = keypoint_module.NormalizedKeypoint _Rect = rect.Rect -_OutputType = interactive_segmenter.InteractiveSegmenterOptions.OutputType _InteractiveSegmenter = interactive_segmenter.InteractiveSegmenter _InteractiveSegmenterOptions = interactive_segmenter.InteractiveSegmenterOptions _RegionOfInterest = interactive_segmenter.RegionOfInterest @@ -200,15 +200,14 @@ class InteractiveSegmenterTest(parameterized.TestCase): raise ValueError('model_file_type is invalid.') options = _InteractiveSegmenterOptions( - base_options=base_options, output_type=_OutputType.CATEGORY_MASK + base_options=base_options, output_category_mask=True ) segmenter = _InteractiveSegmenter.create_from_options(options) # Performs image segmentation on the input. roi = _RegionOfInterest(format=roi_format, keypoint=keypoint) - category_masks = segmenter.segment(self.test_image, roi) - self.assertLen(category_masks, 1) - category_mask = category_masks[0] + segmentation_result = segmenter.segment(self.test_image, roi) + category_mask = segmentation_result.category_mask result_pixels = category_mask.numpy_view().flatten() # Check if data type of `category_mask` is correct. @@ -219,7 +218,7 @@ class InteractiveSegmenterTest(parameterized.TestCase): self.assertTrue( _similar_to_uint8_mask( - category_masks[0], test_seg_image, similarity_threshold + category_mask, test_seg_image, similarity_threshold ), ( 'Number of pixels in the candidate mask differing from that of the' @@ -253,13 +252,12 @@ 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, output_type=_OutputType.CONFIDENCE_MASK - ) + options = _InteractiveSegmenterOptions(base_options=base_options) with _InteractiveSegmenter.create_from_options(options) as segmenter: # Perform segmentation - confidence_masks = segmenter.segment(self.test_image, roi) + segmentation_result = segmenter.segment(self.test_image, roi) + confidence_masks = segmentation_result.confidence_masks # Check if confidence mask shape is correct. self.assertLen( @@ -286,16 +284,15 @@ class InteractiveSegmenterTest(parameterized.TestCase): ) # Run segmentation on the model in CONFIDENCE_MASK mode. - options = _InteractiveSegmenterOptions( - base_options=base_options, output_type=_OutputType.CONFIDENCE_MASK - ) + options = _InteractiveSegmenterOptions(base_options=base_options) with _InteractiveSegmenter.create_from_options(options) as segmenter: # Perform segmentation image_processing_options = _ImageProcessingOptions(rotation_degrees=-90) - confidence_masks = segmenter.segment( + segmentation_result = segmenter.segment( self.test_image, roi, image_processing_options ) + confidence_masks = segmentation_result.confidence_masks # Check if confidence mask shape is correct. self.assertLen( @@ -313,9 +310,7 @@ class InteractiveSegmenterTest(parameterized.TestCase): ) # Run segmentation on the model in CONFIDENCE_MASK mode. - options = _InteractiveSegmenterOptions( - base_options=base_options, output_type=_OutputType.CONFIDENCE_MASK - ) + options = _InteractiveSegmenterOptions(base_options=base_options) with self.assertRaisesRegex( ValueError, "This task doesn't support region-of-interest." diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index e50ffbf7..10277317 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -31,7 +31,6 @@ from mediapipe.tasks.python.vision.core import base_vision_task_api from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module from mediapipe.tasks.python.vision.core import vision_task_running_mode -ImageSegmenterResult = List[image_module.Image] _NormalizedRect = rect.NormalizedRect _BaseOptions = base_options_module.BaseOptions _SegmenterOptionsProto = segmenter_options_pb2.SegmenterOptions @@ -42,8 +41,10 @@ _RunningMode = vision_task_running_mode.VisionTaskRunningMode _ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions _TaskInfo = task_info_module.TaskInfo -_SEGMENTATION_OUT_STREAM_NAME = 'segmented_mask_out' -_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' +_CONFIDENCE_MASKS_STREAM_NAME = 'confidence_masks' +_CONFIDENCE_MASKS_TAG = 'CONFIDENCE_MASKS' +_CATEGORY_MASK_STREAM_NAME = 'category_mask' +_CATEGORY_MASK_TAG = 'CATEGORY_MASK' _IMAGE_IN_STREAM_NAME = 'image_in' _IMAGE_OUT_STREAM_NAME = 'image_out' _IMAGE_TAG = 'IMAGE' @@ -53,6 +54,12 @@ _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph' _MICRO_SECONDS_PER_MILLISECOND = 1000 +@dataclasses.dataclass +class ImageSegmenterResult: + confidence_masks: Optional[List[image_module.Image]] = None + category_mask: Optional[image_module.Image] = None + + @dataclasses.dataclass class ImageSegmenterOptions: """Options for the image segmenter task. @@ -64,19 +71,13 @@ class ImageSegmenterOptions: 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. - 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_confidence_masks: Whether to output confidence masks. + output_category_mask: Whether to output category mask. 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. """ - class OutputType(enum.Enum): - UNSPECIFIED = 0 - CATEGORY_MASK = 1 - CONFIDENCE_MASK = 2 - class Activation(enum.Enum): NONE = 0 SIGMOID = 1 @@ -84,7 +85,8 @@ class ImageSegmenterOptions: base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE - output_type: Optional[OutputType] = OutputType.CATEGORY_MASK + output_confidence_masks: bool = True + output_category_mask: bool = False activation: Optional[Activation] = Activation.NONE result_callback: Optional[ Callable[[ImageSegmenterResult, image_module.Image, int], None] @@ -98,7 +100,7 @@ class ImageSegmenterOptions: False if self.running_mode == _RunningMode.IMAGE else True ) segmenter_options_proto = _SegmenterOptionsProto( - output_type=self.output_type.value, activation=self.activation.value + activation=self.activation.value ) return _ImageSegmenterGraphOptionsProto( base_options=base_options_proto, @@ -177,27 +179,48 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): 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( - output_packets[_SEGMENTATION_OUT_STREAM_NAME] - ) + + segmentation_result = ImageSegmenterResult() + + if options.output_confidence_masks: + segmentation_result.confidence_masks = packet_getter.get_image_list( + output_packets[_CONFIDENCE_MASKS_STREAM_NAME] + ) + + if options.output_category_mask: + segmentation_result.category_mask = packet_getter.get_image( + output_packets[_CATEGORY_MASK_STREAM_NAME] + ) + image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME]) - timestamp = output_packets[_SEGMENTATION_OUT_STREAM_NAME].timestamp + timestamp = output_packets[_IMAGE_OUT_STREAM_NAME].timestamp options.result_callback( segmentation_result, image, timestamp.value // _MICRO_SECONDS_PER_MILLISECOND, ) + output_streams = [ + ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]), + ] + + if options.output_confidence_masks: + output_streams.append( + ':'.join([_CONFIDENCE_MASKS_TAG, _CONFIDENCE_MASKS_STREAM_NAME]) + ) + + if options.output_category_mask: + output_streams.append( + ':'.join([_CATEGORY_MASK_TAG, _CATEGORY_MASK_STREAM_NAME]) + ) + task_info = _TaskInfo( task_graph=_TASK_GRAPH_NAME, input_streams=[ ':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME]), ':'.join([_NORM_RECT_TAG, _NORM_RECT_STREAM_NAME]), ], - output_streams=[ - ':'.join([_SEGMENTATION_TAG, _SEGMENTATION_OUT_STREAM_NAME]), - ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]), - ], + output_streams=output_streams, task_options=options, ) return cls( @@ -240,9 +263,18 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): normalized_rect.to_pb2() ), }) - segmentation_result = packet_getter.get_image_list( - output_packets[_SEGMENTATION_OUT_STREAM_NAME] - ) + segmentation_result = ImageSegmenterResult() + + if _CONFIDENCE_MASKS_STREAM_NAME in output_packets: + segmentation_result.confidence_masks = packet_getter.get_image_list( + output_packets[_CONFIDENCE_MASKS_STREAM_NAME] + ) + + if _CATEGORY_MASK_STREAM_NAME in output_packets: + segmentation_result.category_mask = packet_getter.get_image( + output_packets[_CATEGORY_MASK_STREAM_NAME] + ) + return segmentation_result def segment_for_video( @@ -285,9 +317,19 @@ class ImageSegmenter(base_vision_task_api.BaseVisionTaskApi): normalized_rect.to_pb2() ).at(timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND), }) - segmentation_result = packet_getter.get_image_list( - output_packets[_SEGMENTATION_OUT_STREAM_NAME] - ) + segmentation_result = ImageSegmenterResult() + + if _CONFIDENCE_MASKS_STREAM_NAME in output_packets: + segmentation_result.confidence_masks = packet_getter.get_image_list( + output_packets[_CONFIDENCE_MASKS_STREAM_NAME] + ) + + if _CATEGORY_MASK_STREAM_NAME in output_packets: + segmentation_result.category_mask = packet_getter.get_image( + output_packets[_CATEGORY_MASK_STREAM_NAME] + ) + + return segmentation_result def segment_async( diff --git a/mediapipe/tasks/python/vision/interactive_segmenter.py b/mediapipe/tasks/python/vision/interactive_segmenter.py index 12a30b6e..6f423d76 100644 --- a/mediapipe/tasks/python/vision/interactive_segmenter.py +++ b/mediapipe/tasks/python/vision/interactive_segmenter.py @@ -41,8 +41,10 @@ _RunningMode = vision_task_running_mode.VisionTaskRunningMode _ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions _TaskInfo = task_info_module.TaskInfo -_SEGMENTATION_OUT_STREAM_NAME = 'segmented_mask_out' -_SEGMENTATION_TAG = 'GROUPED_SEGMENTATION' +_CONFIDENCE_MASKS_STREAM_NAME = 'confidence_masks' +_CONFIDENCE_MASKS_TAG = 'CONFIDENCE_MASKS' +_CATEGORY_MASK_STREAM_NAME = 'category_mask' +_CATEGORY_MASK_TAG = 'CATEGORY_MASK' _IMAGE_IN_STREAM_NAME = 'image_in' _IMAGE_OUT_STREAM_NAME = 'image_out' _ROI_STREAM_NAME = 'roi_in' @@ -55,32 +57,32 @@ _TASK_GRAPH_NAME = ( ) +@dataclasses.dataclass +class InteractiveSegmenterResult: + confidence_masks: Optional[List[image_module.Image]] = None + category_mask: Optional[image_module.Image] = None + + @dataclasses.dataclass class InteractiveSegmenterOptions: """Options for the interactive segmenter task. Attributes: base_options: Base options for the interactive segmenter task. - output_type: The output mask type allows specifying the type of - post-processing to perform on the raw model results. + output_confidence_masks: Whether to output confidence masks. + output_category_mask: Whether to output category mask. """ - class OutputType(enum.Enum): - UNSPECIFIED = 0 - CATEGORY_MASK = 1 - CONFIDENCE_MASK = 2 - base_options: _BaseOptions - output_type: Optional[OutputType] = OutputType.CATEGORY_MASK + output_confidence_masks: bool = True + output_category_mask: bool = False @doc_controls.do_not_generate_docs def to_pb2(self) -> _ImageSegmenterGraphOptionsProto: """Generates an InteractiveSegmenterOptions protobuf object.""" base_options_proto = self.base_options.to_pb2() base_options_proto.use_stream_mode = False - segmenter_options_proto = _SegmenterOptionsProto( - output_type=self.output_type.value - ) + segmenter_options_proto = _SegmenterOptionsProto() return _ImageSegmenterGraphOptionsProto( base_options=base_options_proto, segmenter_options=segmenter_options_proto, @@ -192,6 +194,20 @@ class InteractiveSegmenter(base_vision_task_api.BaseVisionTaskApi): RuntimeError: If other types of error occurred. """ + output_streams = [ + ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]), + ] + + if options.output_confidence_masks: + output_streams.append( + ':'.join([_CONFIDENCE_MASKS_TAG, _CONFIDENCE_MASKS_STREAM_NAME]) + ) + + if options.output_category_mask: + output_streams.append( + ':'.join([_CATEGORY_MASK_TAG, _CATEGORY_MASK_STREAM_NAME]) + ) + task_info = _TaskInfo( task_graph=_TASK_GRAPH_NAME, input_streams=[ @@ -199,10 +215,7 @@ class InteractiveSegmenter(base_vision_task_api.BaseVisionTaskApi): ':'.join([_ROI_TAG, _ROI_STREAM_NAME]), ':'.join([_NORM_RECT_TAG, _NORM_RECT_STREAM_NAME]), ], - output_streams=[ - ':'.join([_SEGMENTATION_TAG, _SEGMENTATION_OUT_STREAM_NAME]), - ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]), - ], + output_streams=output_streams, task_options=options, ) return cls( @@ -216,7 +229,7 @@ class InteractiveSegmenter(base_vision_task_api.BaseVisionTaskApi): image: image_module.Image, roi: RegionOfInterest, image_processing_options: Optional[_ImageProcessingOptions] = None, - ) -> List[image_module.Image]: + ) -> InteractiveSegmenterResult: """Performs the actual segmentation task on the provided MediaPipe Image. The image can be of any size with format RGB. @@ -248,7 +261,16 @@ class InteractiveSegmenter(base_vision_task_api.BaseVisionTaskApi): normalized_rect.to_pb2() ), }) - segmentation_result = packet_getter.get_image_list( - output_packets[_SEGMENTATION_OUT_STREAM_NAME] - ) + segmentation_result = InteractiveSegmenterResult() + + if _CONFIDENCE_MASKS_STREAM_NAME in output_packets: + segmentation_result.confidence_masks = packet_getter.get_image_list( + output_packets[_CONFIDENCE_MASKS_STREAM_NAME] + ) + + if _CATEGORY_MASK_STREAM_NAME in output_packets: + segmentation_result.category_mask = packet_getter.get_image( + output_packets[_CATEGORY_MASK_STREAM_NAME] + ) + return segmentation_result diff --git a/mediapipe/tasks/testdata/vision/BUILD b/mediapipe/tasks/testdata/vision/BUILD index 0de0c255..23dc3ed5 100644 --- a/mediapipe/tasks/testdata/vision/BUILD +++ b/mediapipe/tasks/testdata/vision/BUILD @@ -77,6 +77,7 @@ mediapipe_files(srcs = [ "portrait_selfie_segmentation_landscape_expected_category_mask.jpg", "pose.jpg", "pose_detection.tflite", + "ptm_512_hdt_ptm_woid.tflite", "pose_landmark_lite.tflite", "pose_landmarker.task", "right_hands.jpg", @@ -187,6 +188,7 @@ filegroup( "mobilenet_v3_small_100_224_embedder.tflite", "palm_detection_full.tflite", "pose_detection.tflite", + "ptm_512_hdt_ptm_woid.tflite", "pose_landmark_lite.tflite", "pose_landmarker.task", "selfie_segm_128_128_3.tflite", From a03fa448dc09d75da7f5eecae16a1a2909a46df6 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 13 Apr 2023 11:55:37 -0700 Subject: [PATCH 2/9] 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 32792519..aa557281 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 6af15aa0..aea4f8a1 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." From a036bf70cc014ce5ba435ab53e3aa25b38a7ad0e Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 13 Apr 2023 21:09:01 -0700 Subject: [PATCH 3/9] Removed Activation from ImageSegmenterOptions --- .../tasks/python/test/vision/image_segmenter_test.py | 4 +--- mediapipe/tasks/python/vision/image_segmenter.py | 10 +--------- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_segmenter_test.py b/mediapipe/tasks/python/test/vision/image_segmenter_test.py index aa557281..de0a263f 100644 --- a/mediapipe/tasks/python/test/vision/image_segmenter_test.py +++ b/mediapipe/tasks/python/test/vision/image_segmenter_test.py @@ -15,7 +15,6 @@ import enum import os -from typing import List from unittest import mock from absl.testing import absltest @@ -34,7 +33,6 @@ ImageSegmenterResult = image_segmenter.ImageSegmenterResult _BaseOptions = base_options_module.BaseOptions _Image = image_module.Image _ImageFormat = image_frame.ImageFormat -_Activation = image_segmenter.ImageSegmenterOptions.Activation _ImageSegmenter = image_segmenter.ImageSegmenter _ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions _RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode @@ -191,7 +189,7 @@ class ImageSegmenterTest(parameterized.TestCase): # Run segmentation on the model in CONFIDENCE_MASK mode. options = _ImageSegmenterOptions( base_options=base_options, output_category_mask=False, - output_confidence_masks=True, activation=_Activation.SOFTMAX + output_confidence_masks=True ) with _ImageSegmenter.create_from_options(options) as segmenter: diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index 10277317..c5204db4 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -78,16 +78,10 @@ class ImageSegmenterOptions: is set to the live stream mode. """ - class Activation(enum.Enum): - NONE = 0 - SIGMOID = 1 - SOFTMAX = 2 - base_options: _BaseOptions running_mode: _RunningMode = _RunningMode.IMAGE output_confidence_masks: bool = True output_category_mask: bool = False - activation: Optional[Activation] = Activation.NONE result_callback: Optional[ Callable[[ImageSegmenterResult, image_module.Image, int], None] ] = None @@ -99,9 +93,7 @@ class ImageSegmenterOptions: base_options_proto.use_stream_mode = ( False if self.running_mode == _RunningMode.IMAGE else True ) - segmenter_options_proto = _SegmenterOptionsProto( - activation=self.activation.value - ) + segmenter_options_proto = _SegmenterOptionsProto() return _ImageSegmenterGraphOptionsProto( base_options=base_options_proto, segmenter_options=segmenter_options_proto, From a745b71f97aa49a5533ddb282bfee2bd30654aa4 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Thu, 13 Apr 2023 21:11:20 -0700 Subject: [PATCH 4/9] Removed unused import --- mediapipe/tasks/python/vision/image_segmenter.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mediapipe/tasks/python/vision/image_segmenter.py b/mediapipe/tasks/python/vision/image_segmenter.py index c5204db4..8edabe32 100644 --- a/mediapipe/tasks/python/vision/image_segmenter.py +++ b/mediapipe/tasks/python/vision/image_segmenter.py @@ -14,7 +14,6 @@ """MediaPipe image segmenter task.""" import dataclasses -import enum from typing import Callable, List, Mapping, Optional from mediapipe.python import packet_creator From b4e27c137e2c5c014c3e849fec7ec0ed92759fbf Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 18 Apr 2023 00:52:09 -0700 Subject: [PATCH 5/9] Internal change PiperOrigin-RevId: 525069421 --- .../objectron/calculators/frame_annotation_tracker.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mediapipe/modules/objectron/calculators/frame_annotation_tracker.cc b/mediapipe/modules/objectron/calculators/frame_annotation_tracker.cc index eebf8857..1685a4f6 100644 --- a/mediapipe/modules/objectron/calculators/frame_annotation_tracker.cc +++ b/mediapipe/modules/objectron/calculators/frame_annotation_tracker.cc @@ -24,8 +24,8 @@ namespace mediapipe { void FrameAnnotationTracker::AddDetectionResult( const FrameAnnotation& frame_annotation) { - const int64 time_us = - static_cast(std::round(frame_annotation.timestamp())); + const int64_t time_us = + static_cast(std::round(frame_annotation.timestamp())); for (const auto& object_annotation : frame_annotation.annotations()) { detected_objects_[time_us + object_annotation.object_id()] = object_annotation; @@ -37,7 +37,7 @@ FrameAnnotation FrameAnnotationTracker::ConsolidateTrackingResult( absl::flat_hash_set* cancel_object_ids) { CHECK(cancel_object_ids != nullptr); FrameAnnotation frame_annotation; - std::vector keys_to_be_deleted; + std::vector keys_to_be_deleted; for (const auto& detected_obj : detected_objects_) { const int object_id = detected_obj.second.object_id(); if (cancel_object_ids->contains(object_id)) { From 88a10de345926cab4432c4ec61339b48af28f87e Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 18 Apr 2023 02:12:28 -0700 Subject: [PATCH 6/9] Internal change PiperOrigin-RevId: 525084368 --- mediapipe/gpu/gl_texture_buffer.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mediapipe/gpu/gl_texture_buffer.cc b/mediapipe/gpu/gl_texture_buffer.cc index 69b9889c..f1497f74 100644 --- a/mediapipe/gpu/gl_texture_buffer.cc +++ b/mediapipe/gpu/gl_texture_buffer.cc @@ -64,7 +64,7 @@ std::unique_ptr GlTextureBuffer::Create( int actual_ws = image_frame.WidthStep(); int alignment = 0; std::unique_ptr temp; - const uint8* data = image_frame.PixelData(); + const uint8_t* data = image_frame.PixelData(); // Let's see if the pixel data is tightly aligned to one of the alignments // supported by OpenGL, preferring 4 if possible since it's the default. From 3e0ed2ced053faae2f078dd2d53052c2f578120f Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 18 Apr 2023 10:10:58 -0700 Subject: [PATCH 7/9] Internal Changes PiperOrigin-RevId: 525180095 --- .../python/vision/object_detector/BUILD | 6 +----- .../object_detector/object_detector_test.py | 17 +++++++---------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/mediapipe/model_maker/python/vision/object_detector/BUILD b/mediapipe/model_maker/python/vision/object_detector/BUILD index f3d4407d..b97d215d 100644 --- a/mediapipe/model_maker/python/vision/object_detector/BUILD +++ b/mediapipe/model_maker/python/vision/object_detector/BUILD @@ -175,11 +175,7 @@ py_test( data = [":testdata"], tags = ["requires-net:external"], deps = [ - ":dataset", - ":hyperparameters", - ":model_spec", - ":object_detector", - ":object_detector_options", + ":object_detector_import", "//mediapipe/tasks/python/test:test_utils", ], ) diff --git a/mediapipe/model_maker/python/vision/object_detector/object_detector_test.py b/mediapipe/model_maker/python/vision/object_detector/object_detector_test.py index df6b58a0..02f773e6 100644 --- a/mediapipe/model_maker/python/vision/object_detector/object_detector_test.py +++ b/mediapipe/model_maker/python/vision/object_detector/object_detector_test.py @@ -19,11 +19,7 @@ from unittest import mock as unittest_mock from absl.testing import parameterized import tensorflow as tf -from mediapipe.model_maker.python.vision.object_detector import dataset -from mediapipe.model_maker.python.vision.object_detector import hyperparameters -from mediapipe.model_maker.python.vision.object_detector import model_spec as ms -from mediapipe.model_maker.python.vision.object_detector import object_detector -from mediapipe.model_maker.python.vision.object_detector import object_detector_options +from mediapipe.model_maker.python.vision import object_detector from mediapipe.tasks.python.test import test_utils as task_test_utils @@ -33,7 +29,7 @@ class ObjectDetectorTest(tf.test.TestCase, parameterized.TestCase): super().setUp() dataset_folder = task_test_utils.get_test_data_path('coco_data') cache_dir = self.create_tempdir() - self.data = dataset.Dataset.from_coco_folder( + self.data = object_detector.Dataset.from_coco_folder( dataset_folder, cache_dir=cache_dir ) # Mock tempfile.gettempdir() to be unique for each test to avoid race @@ -48,15 +44,16 @@ class ObjectDetectorTest(tf.test.TestCase, parameterized.TestCase): self.addCleanup(mock_gettempdir.stop) def test_object_detector(self): - hparams = hyperparameters.HParams( + hparams = object_detector.HParams( epochs=1, batch_size=2, learning_rate=0.9, shuffle=False, export_dir=self.create_tempdir(), ) - options = object_detector_options.ObjectDetectorOptions( - supported_model=ms.SupportedModels.MOBILENET_V2, hparams=hparams + options = object_detector.ObjectDetectorOptions( + supported_model=object_detector.SupportedModels.MOBILENET_V2, + hparams=hparams, ) # Test `create`` model = object_detector.ObjectDetector.create( @@ -79,7 +76,7 @@ class ObjectDetectorTest(tf.test.TestCase, parameterized.TestCase): self.assertGreater(os.path.getsize(output_metadata_file), 0) # Test `quantization_aware_training` - qat_hparams = hyperparameters.QATHParams( + qat_hparams = object_detector.QATHParams( learning_rate=0.9, batch_size=2, epochs=1, From 64d1e74c208d04b999998dbb2560adc266280a7e Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 18 Apr 2023 10:18:19 -0700 Subject: [PATCH 8/9] Internal MediaPipe Tasks change PiperOrigin-RevId: 525182282 --- .../python/test/text/text_embedder_test.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/mediapipe/tasks/python/test/text/text_embedder_test.py b/mediapipe/tasks/python/test/text/text_embedder_test.py index 78e98a1b..62d162f6 100644 --- a/mediapipe/tasks/python/test/text/text_embedder_test.py +++ b/mediapipe/tasks/python/test/text/text_embedder_test.py @@ -32,6 +32,7 @@ _TextEmbedderOptions = text_embedder.TextEmbedderOptions _BERT_MODEL_FILE = 'mobilebert_embedding_with_metadata.tflite' _REGEX_MODEL_FILE = 'regex_one_embedding_with_metadata.tflite' +_USE_MODEL_FILE = 'universal_sentence_encoder_qa_with_metadata.tflite' _TEST_DATA_DIR = 'mediapipe/tasks/testdata/text' # Tolerance for embedding vector coordinate values. _EPSILON = 1e-4 @@ -138,6 +139,24 @@ class TextEmbedderTest(parameterized.TestCase): 16, (0.549632, 0.552879), ), + ( + False, + False, + _USE_MODEL_FILE, + ModelFileType.FILE_NAME, + 0.851961, + 100, + (1.422951, 1.404664), + ), + ( + True, + False, + _USE_MODEL_FILE, + ModelFileType.FILE_CONTENT, + 0.851961, + 100, + (0.127049, 0.125416), + ), ) def test_embed(self, l2_normalize, quantize, model_name, model_file_type, expected_similarity, expected_size, expected_first_values): @@ -213,6 +232,24 @@ class TextEmbedderTest(parameterized.TestCase): 16, (0.549632, 0.552879), ), + ( + False, + False, + _USE_MODEL_FILE, + ModelFileType.FILE_NAME, + 0.851961, + 100, + (1.422951, 1.404664), + ), + ( + True, + False, + _USE_MODEL_FILE, + ModelFileType.FILE_CONTENT, + 0.851961, + 100, + (0.127049, 0.125416), + ), ) def test_embed_in_context(self, l2_normalize, quantize, model_name, model_file_type, expected_similarity, expected_size, @@ -251,6 +288,7 @@ class TextEmbedderTest(parameterized.TestCase): @parameterized.parameters( # TODO: The similarity should likely be lower (_BERT_MODEL_FILE, 0.980880), + (_USE_MODEL_FILE, 0.780334), ) def test_embed_with_different_themes(self, model_file, expected_similarity): # Creates embedder. From d7039c90dc9eebf047d099ee97c980c23c0c51ad Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Tue, 18 Apr 2023 14:02:43 -0700 Subject: [PATCH 9/9] Update WASM for Alpha 11 PiperOrigin-RevId: 525245471 --- third_party/wasm_files.bzl | 48 +++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/third_party/wasm_files.bzl b/third_party/wasm_files.bzl index 148b5970..a484d2f8 100644 --- a/third_party/wasm_files.bzl +++ b/third_party/wasm_files.bzl @@ -12,72 +12,72 @@ def wasm_files(): http_file( name = "com_google_mediapipe_wasm_audio_wasm_internal_js", - sha256 = "0eca68e2291a548b734bcab5db4c9e6b997e852ea7e19228003b9e2a78c7c646", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_internal.js?generation=1681328323089931"], + sha256 = "b810de53d7ccf991b9c70fcdf7e88b5c3f2942ae766436f22be48159b6a7e687", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_internal.js?generation=1681849488227617"], ) http_file( name = "com_google_mediapipe_wasm_audio_wasm_internal_wasm", - sha256 = "69bc95af5b783b510ec1842d6fb9594254907d8e1334799c5753164878a7dcac", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_internal.wasm?generation=1681328325829340"], + sha256 = "26d91147e5c6c8a92e0a4ebf59599068a3cff6108847b793ef33ac23e98eddb9", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_internal.wasm?generation=1681849491546937"], ) http_file( name = "com_google_mediapipe_wasm_audio_wasm_nosimd_internal_js", - sha256 = "88a0176cc80d6a1eb175a5105df705cf8b8684cf13f6db0a264af0b67b65a22a", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_nosimd_internal.js?generation=1681328328330829"], + sha256 = "b38e37b3024692558eaaba159921fedd3297d1a09bba1c16a06fed327845b0bd", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_nosimd_internal.js?generation=1681849494099698"], ) http_file( name = "com_google_mediapipe_wasm_audio_wasm_nosimd_internal_wasm", - sha256 = "1cc0c3db7d252801be4b090d8bbba61f308cc3dd5efe197319581d3af29495c7", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_nosimd_internal.wasm?generation=1681328331085637"], + sha256 = "6a8e73d2e926565046e16adf1748f0f8ec5135fafe7eb8b9c83892e64c1a449a", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/audio_wasm_nosimd_internal.wasm?generation=1681849496451970"], ) http_file( name = "com_google_mediapipe_wasm_text_wasm_internal_js", - sha256 = "d9cd100b6d330d36f7749fe5fc64a2cdd0abb947a0376e6140784cfb0361a4e2", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_internal.js?generation=1681328333442454"], + sha256 = "785cba67b623b1dc66dc3621e97fd6b30edccbb408184a3094d0aa68ddd5becb", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_internal.js?generation=1681849498746265"], ) http_file( name = "com_google_mediapipe_wasm_text_wasm_internal_wasm", - sha256 = "30a2fcca630bdad6e99173ea7d0d8c5d7086aedf393d0159fa05bf9d08d4ff65", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_internal.wasm?generation=1681328335803336"], + sha256 = "a858b8a2e8b40e9c936b66566c5aefd396536c4e936459ab9ae7e239621adc14", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_internal.wasm?generation=1681849501370461"], ) http_file( name = "com_google_mediapipe_wasm_text_wasm_nosimd_internal_js", - sha256 = "70ca2bd15c56e0ce7bb10ff2188b4a1f9eafbb657eb9424e4cab8d7b29179871", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_nosimd_internal.js?generation=1681328338162884"], + sha256 = "5292f1442d5e5c037e7cffb78a8c2d71255348ca2c3bd759b314bdbedd5590c2", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_nosimd_internal.js?generation=1681849503379116"], ) http_file( name = "com_google_mediapipe_wasm_text_wasm_nosimd_internal_wasm", - sha256 = "8221b385905f36a769d7731a0adbe18b681bcb873561890429ca84278c67c3fd", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_nosimd_internal.wasm?generation=1681328340808115"], + sha256 = "e44b48ab29ee1d8befec804e9a63445c56266b679d19fb476d556ca621f0e493", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/text_wasm_nosimd_internal.wasm?generation=1681849505997020"], ) http_file( name = "com_google_mediapipe_wasm_vision_wasm_internal_js", - sha256 = "07692acd8202adafebd35dbcd7e2b8e88a76d4a0e6b9229cb3cad59503eeddc7", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_internal.js?generation=1681328343147709"], + sha256 = "205855eba70464a92b9d00e90acac15c51a9f76192f900e697304ac6dea8f714", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_internal.js?generation=1681849508414277"], ) http_file( name = "com_google_mediapipe_wasm_vision_wasm_internal_wasm", - sha256 = "03bf553fa6a768b0d70103a5e7d835b6b37371ff44e201c3392f22e0879737c3", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_internal.wasm?generation=1681328345605574"], + sha256 = "c0cbd0df3adb2a9cd1331d14f522d2bae9f8adc9f1b35f92cbbc4b782b190cef", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_internal.wasm?generation=1681849510936608"], ) http_file( name = "com_google_mediapipe_wasm_vision_wasm_nosimd_internal_js", - sha256 = "36697be14f921985eac15d1447ec8a260817b05ade1c9bb3ca7e906e0f047ec0", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_nosimd_internal.js?generation=1681328348025082"], + sha256 = "0969812de4d3573198fa2eba4f5b0a7e97e98f97bd4215d876543f4925e57b84", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_nosimd_internal.js?generation=1681849513292639"], ) http_file( name = "com_google_mediapipe_wasm_vision_wasm_nosimd_internal_wasm", - sha256 = "103fb145438d61cfecb2e8db3f06b43a5d77a7e3fcea940437fe272227cf2592", - urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_nosimd_internal.wasm?generation=1681328350709881"], + sha256 = "f2ab62c3f8dabab0a573dadf5c105ff81a03c29c70f091f8cf273ae030c0a86f", + urls = ["https://storage.googleapis.com/mediapipe-assets/wasm/vision_wasm_nosimd_internal.wasm?generation=1681849515999000"], )