mediapipe-nightly 0.10.21.post20241223__cp312-cp312-manylinux_2_28_x86_64.whl

mediapipe/tasks/python/metadata/metadata_displayer_cli.py

@@ -0,0 +1,34 @@
+# Copyright 2022 The MediaPipe Authors.
+#
+# 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.
+# ==============================================================================
+"""CLI tool for display metadata."""
+
+from absl import app
+from absl import flags
+
+from mediapipe.tasks.python.metadata import metadata
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('model_path', None, 'Path to the TFLite model file.')
+flags.DEFINE_string('export_json_path', None, 'Path to the output JSON file.')
+
+
+def main(_):
+  displayer = metadata.MetadataDisplayer.with_model_file(FLAGS.model_path)
+  with open(FLAGS.export_json_path, 'w') as f:
+    f.write(displayer.get_metadata_json())
+
+
+if __name__ == '__main__':
+  app.run(main)

mediapipe/tasks/python/metadata/metadata_writers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2022 The MediaPipe Authors.
+#
+# 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/tasks/python/metadata/metadata_writers/face_stylizer.py

@@ -0,0 +1,138 @@
+# Copyright 2023 The MediaPipe Authors.
+#
+# 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.
+# ==============================================================================
+"""Writes metadata and creates model asset bundle for face stylizer."""
+
+import os
+import tempfile
+from typing import List
+
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
+from mediapipe.tasks.python.metadata.metadata_writers import model_asset_bundle_utils
+
+_MODEL_NAME = "FaceStylizer"
+_MODEL_DESCRIPTION = "Performs face stylization on images."
+_FACE_DETECTOR_MODEL = "face_detector.tflite"
+_FACE_LANDMARKS_DETECTOR_MODEL = "face_landmarks_detector.tflite"
+_FACE_STYLIZER_MODEL = "face_stylizer.tflite"
+_FACE_STYLIZER_TASK = "face_stylizer.task"
+
+
+class MetadataWriter:
+  """MetadataWriter to write the metadata for face stylizer."""
+
+  def __init__(
+      self,
+      face_detector_model_buffer: bytearray,
+      face_landmarks_detector_model_buffer: bytearray,
+      face_stylizer_metadata_writer: metadata_writer.MetadataWriter,
+  ) -> None:
+    """Initializes MetadataWriter to write the metadata and model asset bundle.
+
+    Args:
+      face_detector_model_buffer: A valid flatbuffer loaded from the face
+        detector TFLite model file with metadata already packed inside.
+      face_landmarks_detector_model_buffer: A valid flatbuffer loaded from the
+        face landmarks detector TFLite model file with metadata already packed
+        inside.
+      face_stylizer_metadata_writer: Metadata writer to write face stylizer
+        metadata into the TFLite file.
+    """
+    self._face_detector_model_buffer = face_detector_model_buffer
+    self._face_landmarks_detector_model_buffer = (
+        face_landmarks_detector_model_buffer
+    )
+    self._face_stylizer_metadata_writer = face_stylizer_metadata_writer
+    self._temp_folder = tempfile.TemporaryDirectory()
+
+  def __del__(self):
+    if os.path.exists(self._temp_folder.name):
+      self._temp_folder.cleanup()
+
+  @classmethod
+  def create(
+      cls,
+      face_stylizer_model_buffer: bytearray,
+      face_detector_model_buffer: bytearray,
+      face_landmarks_detector_model_buffer: bytearray,
+      input_norm_mean: List[float],
+      input_norm_std: List[float],
+  ) -> "MetadataWriter":
+    """Creates MetadataWriter to write the metadata for face stylizer.
+
+    The parameters required in this method are mandatory when using MediaPipe
+    Tasks.
+
+    Note that only the output TFLite is used for deployment. The output JSON
+    content is used to interpret the metadata content.
+
+    Args:
+      face_stylizer_model_buffer: A valid flatbuffer loaded from the face
+        stylizer TFLite model file.
+      face_detector_model_buffer: A valid flatbuffer loaded from the face
+        detector TFLite model file with metadata already packed inside.
+      face_landmarks_detector_model_buffer: A valid flatbuffer loaded from the
+        face landmarks detector TFLite model file with metadata already packed
+        inside.
+      input_norm_mean: the mean value used in the input tensor normalization for
+        face stylizer model [1].
+      input_norm_std: the std value used in the input tensor normalizarion for
+        face stylizer model [1].
+
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
+
+    Returns:
+      A MetadataWriter object.
+    """
+    face_stylizer_writer = metadata_writer.MetadataWriter(
+        face_stylizer_model_buffer
+    )
+    face_stylizer_writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
+    face_stylizer_writer.add_image_input(input_norm_mean, input_norm_std)
+    return cls(
+        face_detector_model_buffer,
+        face_landmarks_detector_model_buffer,
+        face_stylizer_writer,
+    )
+
+  def populate(self):
+    """Populates the metadata and creates model asset bundle.
+
+    Note that only the output model asset bundle is used for deployment.
+    The output JSON content is used to interpret the face stylizer metadata
+    content.
+
+    Returns:
+      A tuple of (model_asset_bundle_in_bytes, metadata_json_content)
+    """
+    # Write metadata into the face stylizer TFLite model.
+    face_stylizer_model_buffer, face_stylizer_metadata_json = (
+        self._face_stylizer_metadata_writer.populate()
+    )
+    # Create the model asset bundle for the face stylizer task.
+    face_stylizer_models = {
+        _FACE_DETECTOR_MODEL: self._face_detector_model_buffer,
+        _FACE_LANDMARKS_DETECTOR_MODEL: (
+            self._face_landmarks_detector_model_buffer
+        ),
+        _FACE_STYLIZER_MODEL: face_stylizer_model_buffer,
+    }
+    output_path = os.path.join(self._temp_folder.name, _FACE_STYLIZER_TASK)
+    model_asset_bundle_utils.create_model_asset_bundle(
+        face_stylizer_models, output_path
+    )
+    with open(output_path, "rb") as f:
+      face_stylizer_model_bundle_buffer = f.read()
+    return face_stylizer_model_bundle_buffer, face_stylizer_metadata_json

mediapipe/tasks/python/metadata/metadata_writers/image_classifier.py

@@ -0,0 +1,71 @@
+# Copyright 2022 The MediaPipe Authors.
+#
+# 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.
+# ==============================================================================
+"""Writes metadata and label file to the image classifier models."""
+
+from typing import List, Optional
+
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
+
+_MODEL_NAME = "ImageClassifier"
+_MODEL_DESCRIPTION = ("Identify the most prominent object in the image from a "
+                      "known set of categories.")
+
+
+class MetadataWriter(metadata_writer.MetadataWriterBase):
+  """MetadataWriter to write the metadata for image classifier."""
+
+  @classmethod
+  def create(
+      cls,
+      model_buffer: bytearray,
+      input_norm_mean: List[float],
+      input_norm_std: List[float],
+      labels: metadata_writer.Labels,
+      score_calibration: Optional[metadata_writer.ScoreCalibration] = None
+  ) -> "MetadataWriter":
+    """Creates MetadataWriter to write the metadata for image classifier.
+
+    The parameters required in this method are mandatory when using MediaPipe
+    Tasks.
+
+    Note that only the output TFLite is used for deployment. The output JSON
+    content is used to interpret the metadata content.
+
+    Args:
+      model_buffer: A valid flatbuffer loaded from the TFLite model file.
+      input_norm_mean: the mean value used in the input tensor normalization
+        [1].
+      input_norm_std: the std value used in the input tensor normalizarion [1].
+      labels: an instance of Labels helper class used in the output
+        classification tensor [2].
+      score_calibration: A container of the score calibration operation [3] in
+        the classification tensor. Optional if the model does not use score
+        calibration.
+
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
+      [3]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456
+
+    Returns:
+      A MetadataWriter object.
+    """
+    writer = metadata_writer.MetadataWriter(model_buffer)
+    writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
+    writer.add_image_input(input_norm_mean, input_norm_std)
+    writer.add_classification_output(labels, score_calibration)
+    return cls(writer)

mediapipe/tasks/python/metadata/metadata_writers/image_segmenter.py

@@ -0,0 +1,170 @@
+# Copyright 2022 The MediaPipe Authors.
+#
+# 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.
+# ==============================================================================
+"""Writes metadata and label file to the image segmenter models."""
+import enum
+from typing import List, Optional
+
+import flatbuffers
+
+from mediapipe.tasks.metadata import image_segmenter_metadata_schema_py_generated as _segmenter_metadata_fb
+from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
+from mediapipe.tasks.python.metadata import metadata
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_info
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
+from mediapipe.tasks.python.metadata.metadata_writers import writer_utils
+
+
+_MODEL_NAME = "ImageSegmenter"
+_MODEL_DESCRIPTION = (
+    "Semantic image segmentation predicts whether each pixel "
+    "of an image is associated with a certain class."
+)
+
+# Metadata Schema file for image segmenter.
+_FLATC_METADATA_SCHEMA_FILE = metadata.get_path_to_datafile(
+    "../../../metadata/image_segmenter_metadata_schema.fbs",
+)
+
+# Metadata name in custom metadata field. The metadata name is used to get
+# image segmenter metadata from SubGraphMetadata.custom_metadata and
+# shouldn't be changed.
+_METADATA_NAME = "SEGMENTER_METADATA"
+
+
+class Activation(enum.Enum):
+  NONE = 0
+  SIGMOID = 1
+  SOFTMAX = 2
+
+
+# Create an individual method for getting the metadata json file, so that it can
+# be used as a standalone util.
+def convert_to_json(metadata_buffer: bytearray) -> str:
+  """Converts the metadata into a json string.
+
+  Args:
+    metadata_buffer: valid metadata buffer in bytes.
+
+  Returns:
+    Metadata in JSON format.
+
+  Raises:
+    ValueError: error occurred when parsing the metadata schema file.
+  """
+  return metadata.convert_to_json(
+      metadata_buffer,
+      custom_metadata_schema={_METADATA_NAME: _FLATC_METADATA_SCHEMA_FILE},
+  )
+
+
+class ImageSegmenterOptionsMd(metadata_info.CustomMetadataMd):
+  """Image segmenter options metadata."""
+
+  _METADATA_FILE_IDENTIFIER = b"V001"
+
+  def __init__(self, activation: Activation) -> None:
+    """Creates an ImageSegmenterOptionsMd object.
+
+    Args:
+      activation: activation function of the output layer in the image
+        segmenter.
+    """
+    self.activation = activation
+    super().__init__(name=_METADATA_NAME)
+
+  def create_metadata(self) -> _metadata_fb.CustomMetadataT:
+    """Creates the image segmenter options metadata.
+
+    Returns:
+      A Flatbuffers Python object of the custom metadata including image
+      segmenter options metadata.
+    """
+    segmenter_options = _segmenter_metadata_fb.ImageSegmenterOptionsT()
+    segmenter_options.activation = self.activation.value
+
+    # Get the image segmenter options flatbuffer.
+    b = flatbuffers.Builder(0)
+    b.Finish(segmenter_options.Pack(b), self._METADATA_FILE_IDENTIFIER)
+    segmenter_options_buf = b.Output()
+
+    # Add the image segmenter options flatbuffer in custom metadata.
+    custom_metadata = _metadata_fb.CustomMetadataT()
+    custom_metadata.name = self.name
+    custom_metadata.data = segmenter_options_buf
+    return custom_metadata
+
+
+class MetadataWriter(metadata_writer.MetadataWriterBase):
+  """MetadataWriter to write the metadata for image segmenter."""
+
+  @classmethod
+  def create(
+      cls,
+      model_buffer: bytearray,
+      input_norm_mean: List[float],
+      input_norm_std: List[float],
+      labels: Optional[metadata_writer.Labels] = None,
+      activation: Optional[Activation] = None,
+  ) -> "MetadataWriter":
+    """Creates MetadataWriter to write the metadata for image segmenter.
+
+    The parameters required in this method are mandatory when using MediaPipe
+    Tasks.
+
+    Example usage:
+      metadata_writer = image_segmenter.Metadatawriter.create(model_buffer, ...)
+      tflite_content, json_content = metadata_writer.populate()
+
+    When calling `populate` function in this class, it returns TfLite content
+    and JSON content. Note that only the output TFLite is used for deployment.
+    The output JSON content is used to interpret the metadata content.
+
+    Args:
+      model_buffer: A valid flatbuffer loaded from the TFLite model file.
+      input_norm_mean: the mean value used in the input tensor normalization
+        [1].
+      input_norm_std: the std value used in the input tensor normalizarion [1].
+      labels: an instance of Labels helper class used in the output category
+        tensor [2].
+      activation: activation function for the output layer.
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L116
+
+    Returns:
+      A MetadataWriter object.
+    """
+    writer = metadata_writer.MetadataWriter(model_buffer)
+    writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
+    writer.add_image_input(input_norm_mean, input_norm_std)
+    if activation is not None:
+      option_md = ImageSegmenterOptionsMd(activation)
+      writer.add_custom_metadata(option_md)
+    num_output_tensors = writer_utils.get_subgraph(model_buffer).OutputsLength()
+    if num_output_tensors == 2:
+      # For image segmenter model with 2 output tensors, the first one is
+      # quality score, and the second one is matting mask.
+      writer.add_feature_output(
+          "quality score", "The quality score of matting result."
+      )
+    writer.add_segmentation_output(labels=labels)
+    return cls(writer)
+
+  def populate(self) -> tuple[bytearray, str]:
+    model_buf, _ = super().populate()
+    metadata_buf = metadata.get_metadata_buffer(model_buf)
+    json_content = convert_to_json(metadata_buf)
+    return model_buf, json_content