From e719f7b4e22dfcacfa7b500451ce8aa272dead18 Mon Sep 17 00:00:00 2001 From: Yuqi Li Date: Tue, 1 Nov 2022 15:00:21 -0700 Subject: [PATCH] Model Maker Gesture Recognizer: add metadata writer and create model bundle. PiperOrigin-RevId: 485426865 --- .../metadata_writers/metadata_info.py | 64 +++++++++++++++++-- .../metadata_writers/metadata_writer.py | 20 +++++- .../metadata/metadata_writers/writer_utils.py | 20 +++++- .../metadata_writers/metadata_info_test.py | 19 ++++++ .../metadata_writers/metadata_writer_test.py | 58 +++++++++++++++++ mediapipe/tasks/testdata/metadata/BUILD | 2 + .../metadata/score_thresholding_meta.json | 14 ++++ third_party/external_files.bzl | 6 ++ 8 files changed, 194 insertions(+), 9 deletions(-) create mode 100644 mediapipe/tasks/testdata/metadata/score_thresholding_meta.json diff --git a/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py b/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py index 07938d863..76b572e86 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/metadata_info.py @@ -197,6 +197,37 @@ class ScoreCalibrationMd: self._FILE_TYPE) +class ScoreThresholdingMd: + """A container for score thresholding [1] metadata information. + + [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 + """ + + def __init__(self, global_score_threshold: float) -> None: + """Creates a ScoreThresholdingMd object. + + Args: + global_score_threshold: The recommended global threshold below which + results are considered low-confidence and should be filtered out. + """ + self._global_score_threshold = global_score_threshold + + def create_metadata(self) -> _metadata_fb.ProcessUnitT: + """Creates the score thresholding metadata based on the information. + + Returns: + A Flatbuffers Python object of the score thresholding metadata. + """ + score_thresholding = _metadata_fb.ProcessUnitT() + score_thresholding.optionsType = ( + _metadata_fb.ProcessUnitOptions.ScoreThresholdingOptions) + options = _metadata_fb.ScoreThresholdingOptionsT() + options.globalScoreThreshold = self._global_score_threshold + score_thresholding.options = options + return score_thresholding + + class TensorMd: """A container for common tensor metadata information. @@ -374,23 +405,29 @@ class ClassificationTensorMd(TensorMd): tensor. score_calibration_md: information of the score calibration operation [2] in the classification tensor. + score_thresholding_md: information of the score thresholding [3] in the + classification tensor. [1]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99 [2]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456 + [3]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 """ # Min and max float values for classification results. _MIN_FLOAT = 0.0 _MAX_FLOAT = 1.0 - def __init__(self, - name: Optional[str] = None, - description: Optional[str] = None, - label_files: Optional[List[LabelFileMd]] = None, - tensor_type: Optional[int] = None, - score_calibration_md: Optional[ScoreCalibrationMd] = None, - tensor_name: Optional[str] = None) -> None: + def __init__( + self, + name: Optional[str] = None, + description: Optional[str] = None, + label_files: Optional[List[LabelFileMd]] = None, + tensor_type: Optional[int] = None, + score_calibration_md: Optional[ScoreCalibrationMd] = None, + tensor_name: Optional[str] = None, + score_thresholding_md: Optional[ScoreThresholdingMd] = None) -> None: """Initializes the instance of ClassificationTensorMd. Args: @@ -404,6 +441,8 @@ class ClassificationTensorMd(TensorMd): tensor_name: name of the corresponding tensor [3] in the TFLite model. It is used to locate the corresponding classification tensor and decide the order of the tensor metadata [4] when populating model metadata. + score_thresholding_md: information of the score thresholding [5] in the + classification tensor. [1]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99 [2]: @@ -412,8 +451,11 @@ class ClassificationTensorMd(TensorMd): https://github.com/tensorflow/tensorflow/blob/cb67fef35567298b40ac166b0581cd8ad68e5a3a/tensorflow/lite/schema/schema.fbs#L1129-L1136 [4]: https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L623-L640 + [5]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 """ self.score_calibration_md = score_calibration_md + self.score_thresholding_md = score_thresholding_md if tensor_type is _schema_fb.TensorType.UINT8: min_values = [_MIN_UINT8] @@ -443,4 +485,12 @@ class ClassificationTensorMd(TensorMd): tensor_metadata.processUnits = [ self.score_calibration_md.create_metadata() ] + if self.score_thresholding_md: + if tensor_metadata.processUnits: + tensor_metadata.processUnits.append( + self.score_thresholding_md.create_metadata()) + else: + tensor_metadata.processUnits = [ + self.score_thresholding_md.create_metadata() + ] return tensor_metadata diff --git a/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py b/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py index 5a2eaba07..3a9f91239 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/metadata_writer.py @@ -70,6 +70,18 @@ class LabelItem: locale: Optional[str] = None +@dataclasses.dataclass +class ScoreThresholding: + """Parameters to performs thresholding on output tensor values [1]. + + Attributes: + global_score_threshold: The recommended global threshold below which results + are considered low-confidence and should be filtered out. [1]: + https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L468 + """ + global_score_threshold: float + + class Labels(object): """Simple container holding classification labels of a particular tensor. @@ -407,6 +419,7 @@ class MetadataWriter(object): self, labels: Optional[Labels] = None, score_calibration: Optional[ScoreCalibration] = None, + score_thresholding: Optional[ScoreThresholding] = None, name: str = _OUTPUT_CLASSIFICATION_NAME, description: str = _OUTPUT_CLASSIFICATION_DESCRIPTION ) -> 'MetadataWriter': @@ -423,6 +436,7 @@ class MetadataWriter(object): Args: labels: an instance of Labels helper class. score_calibration: an instance of ScoreCalibration helper class. + score_thresholding: an instance of ScoreThresholding. name: Metadata name of the tensor. Note that this is different from tensor name in the flatbuffer. description: human readable description of what the output is. @@ -437,6 +451,10 @@ class MetadataWriter(object): default_score=score_calibration.default_score, file_path=self._export_calibration_file('score_calibration.txt', score_calibration.parameters)) + score_thresholding_md = None + if score_thresholding: + score_thresholding_md = metadata_info.ScoreThresholdingMd( + score_thresholding.global_score_threshold) label_files = None if labels: @@ -453,6 +471,7 @@ class MetadataWriter(object): label_files=label_files, tensor_type=self._output_tensor_type(len(self._output_mds)), score_calibration_md=calibration_md, + score_thresholding_md=score_thresholding_md, ) self._output_mds.append(output_md) return self @@ -545,4 +564,3 @@ class MetadataWriterBase: A tuple of (model_with_metadata_in_bytes, metdata_json_content) """ return self.writer.populate() - diff --git a/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py b/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py index 0a054812b..eff5f553e 100644 --- a/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py +++ b/mediapipe/tasks/python/metadata/metadata_writers/writer_utils.py @@ -14,7 +14,8 @@ # ============================================================================== """Helper methods for writing metadata into TFLite models.""" -from typing import List +from typing import Dict, List +import zipfile from mediapipe.tasks.metadata import schema_py_generated as _schema_fb @@ -83,3 +84,20 @@ def get_subgraph(model_buffer: bytearray) -> _schema_fb.SubGraph: # multiple subgraphs yet, but models with mini-benchmark may have multiple # subgraphs for acceleration evaluation purpose. return model.Subgraphs(0) + + +def create_model_asset_bundle(input_models: Dict[str, bytes], + output_path: str) -> None: + """Creates the model asset bundle. + + Args: + input_models: A dict of input models with key as the model file name and + value as the model content. + output_path: The output file path to save the model asset bundle. + """ + if not input_models or len(input_models) < 2: + raise ValueError("Needs at least two input models for model asset bundle.") + + with zipfile.ZipFile(output_path, mode="w") as zf: + for file_name, file_buffer in input_models.items(): + zf.writestr(file_name, file_buffer) diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py index 75602c83c..33e162607 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_info_test.py @@ -307,6 +307,25 @@ class ScoreCalibrationMdTest(absltest.TestCase): malformed_calibration_file) +class ScoreThresholdingMdTest(absltest.TestCase): + _DEFAULT_GLOBAL_THRESHOLD = 0.5 + _EXPECTED_TENSOR_JSON = test_utils.get_test_data_path( + "score_thresholding_meta.json") + + def test_create_metadata_should_succeed(self): + score_thresholding_md = metadata_info.ScoreThresholdingMd( + global_score_threshold=self._DEFAULT_GLOBAL_THRESHOLD) + + score_thresholding_metadata = score_thresholding_md.create_metadata() + + metadata_json = _metadata.convert_to_json( + _create_dummy_model_metadata_with_process_uint( + score_thresholding_metadata)) + with open(self._EXPECTED_TENSOR_JSON, "r") as f: + expected_json = f.read() + self.assertEqual(metadata_json, expected_json) + + def _create_dummy_model_metadata_with_tensor( tensor_metadata: _metadata_fb.TensorMetadataT) -> bytes: # Create a dummy model using the tensor metadata. diff --git a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py index 51b043c7d..c59f19519 100644 --- a/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py +++ b/mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py @@ -405,6 +405,64 @@ class MetadataWriterForTaskTest(absltest.TestCase): } """) + def test_add_classification_output_with_score_thresholding(self): + writer = metadata_writer.MetadataWriter.create( + self.image_classifier_model_buffer) + writer.add_classification_output( + labels=metadata_writer.Labels().add(['a', 'b', 'c']), + score_thresholding=metadata_writer.ScoreThresholding( + global_score_threshold=0.5)) + _, metadata_json = writer.populate() + print(metadata_json) + self.assertJsonEqual( + metadata_json, """{ + "subgraph_metadata": [ + { + "input_tensor_metadata": [ + { + "name": "input" + } + ], + "output_tensor_metadata": [ + { + "name": "score", + "description": "Score of the labels respectively.", + "content": { + "content_properties_type": "FeatureProperties", + "content_properties": { + } + }, + "process_units": [ + { + "options_type": "ScoreThresholdingOptions", + "options": { + "global_score_threshold": 0.5 + } + } + ], + "stats": { + "max": [ + 1.0 + ], + "min": [ + 0.0 + ] + }, + "associated_files": [ + { + "name": "labels.txt", + "description": "Labels for categories that the model can recognize.", + "type": "TENSOR_AXIS_LABELS" + } + ] + } + ] + } + ], + "min_parser_version": "1.0.0" + } + """) + if __name__ == '__main__': absltest.main() diff --git a/mediapipe/tasks/testdata/metadata/BUILD b/mediapipe/tasks/testdata/metadata/BUILD index 6d7bbab6a..8ed6e1caa 100644 --- a/mediapipe/tasks/testdata/metadata/BUILD +++ b/mediapipe/tasks/testdata/metadata/BUILD @@ -50,6 +50,7 @@ exports_files([ "score_calibration.txt", "score_calibration_file_meta.json", "score_calibration_tensor_meta.json", + "score_thresholding_meta.json", "labels.txt", "mobilenet_v2_1.0_224.json", "mobilenet_v2_1.0_224_quant.json", @@ -91,5 +92,6 @@ filegroup( "score_calibration.txt", "score_calibration_file_meta.json", "score_calibration_tensor_meta.json", + "score_thresholding_meta.json", ], ) diff --git a/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json b/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json new file mode 100644 index 000000000..d67a1aae8 --- /dev/null +++ b/mediapipe/tasks/testdata/metadata/score_thresholding_meta.json @@ -0,0 +1,14 @@ +{ + "subgraph_metadata": [ + { + "input_process_units": [ + { + "options_type": "ScoreThresholdingOptions", + "options": { + "global_score_threshold": 0.5 + } + } + ] + } + ] +} diff --git a/third_party/external_files.bzl b/third_party/external_files.bzl index 4dcbc3bd9..e47dc9812 100644 --- a/third_party/external_files.bzl +++ b/third_party/external_files.bzl @@ -694,6 +694,12 @@ def external_files(): urls = ["https://storage.googleapis.com/mediapipe-assets/score_calibration.txt?generation=1665422847392804"], ) + http_file( + name = "com_google_mediapipe_score_thresholding_meta_json", + sha256 = "7bb74f21c2d7f0237675ed7c09d7b7afd3507c8373f51dc75fa0507852f6ee19", + urls = ["https://storage.googleapis.com/mediapipe-assets/score_thresholding_meta.json?generation=1667273953630766"], + ) + http_file( name = "com_google_mediapipe_segmentation_golden_rotation0_png", sha256 = "9ee993919b753118928ba2d14f7c5c83a6cfc23355e6943dac4ad81eedd73069",