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

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

@@ -0,0 +1,71 @@
+# 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.
+# ==============================================================================
+"""Utility methods for creating the model asset bundles."""
+
+from typing import Dict
+import zipfile
+
+# Alignment that ensures that all uncompressed files in the model bundle file
+# are aligned relative to the start of the file. This lets the files be
+# accessed directly via mmap.
+_ALIGNMENT = 4
+
+
+class AlignZipFile(zipfile.ZipFile):
+  """ZipFile that stores uncompressed files at particular alignment."""
+
+  def __init__(self, *args, alignment: int, **kwargs) -> None:
+    super().__init__(*args, **kwargs)
+    assert alignment > 0
+    self._alignment = alignment
+
+  def _writecheck(self, zinfo: zipfile.ZipInfo) -> None:
+    # Aligned the uncompressed files.
+    if zinfo.compress_type == zipfile.ZIP_STORED:
+      offset = self.fp.tell()
+      header_length = len(zinfo.FileHeader())
+      padding_length = (
+          self._alignment - (offset + header_length) % self._alignment
+      )
+      if padding_length:
+        offset += padding_length
+        self.fp.write(b"\x00" * padding_length)
+        assert self.fp.tell() == offset
+        zinfo.header_offset = offset
+    else:
+      raise ValueError(
+          "Only support the uncompressed file (compress_type =="
+          " zipfile.ZIP_STORED) in zip. The current file compress type is "
+          + str(zinfo.compress_type)
+      )
+    super()._writecheck(zinfo)
+
+
+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 AlignZipFile(output_path, mode="w", alignment=_ALIGNMENT) as zf:
+    for file_name, file_buffer in input_models.items():
+      zf.writestr(file_name, file_buffer)

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

@@ -0,0 +1,331 @@
+# 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 label file to the Object Detector models."""
+
+import dataclasses
+from typing import List, Optional
+
+import flatbuffers
+
+from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
+from mediapipe.tasks.metadata import object_detector_metadata_schema_py_generated as _detector_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
+
+_MODEL_NAME = "ObjectDetector"
+_MODEL_DESCRIPTION = (
+    "Identify which of a known set of objects might be present and provide "
+    "information about their positions within the given image or a video "
+    "stream."
+)
+
+# Metadata Schema file for object detector.
+_FLATC_METADATA_SCHEMA_FILE = metadata.get_path_to_datafile(
+    "../../../metadata/object_detector_metadata_schema.fbs",
+)
+
+# Metadata name in custom metadata field. The metadata name is used to get
+# object detector metadata from SubGraphMetadata.custom_metadata and shouldn't
+# be changed.
+_METADATA_NAME = "DETECTOR_METADATA"
+
+
+@dataclasses.dataclass
+class FixedAnchor:
+  """A fixed size anchor."""
+
+  x_center: float
+  y_center: float
+  width: Optional[float]
+  height: Optional[float]
+
+
+@dataclasses.dataclass
+class FixedAnchorsSchema:
+  """The schema for a list of anchors with fixed size."""
+
+  anchors: List[FixedAnchor]
+
+
+@dataclasses.dataclass
+class SsdAnchorsOptions:
+  """The ssd anchors options used in object detector model."""
+
+  fixed_anchors_schema: Optional[FixedAnchorsSchema]
+
+
+@dataclasses.dataclass
+class TensorsDecodingOptions:
+  """The decoding options to convert model output tensors to detections."""
+
+  # The number of output classes predicted by the detection model.
+  num_classes: int
+  # The number of output boxes predicted by the detection model.
+  num_boxes: int
+  # The number of output values per boxes predicted by the detection
+  # model. The values contain bounding boxes, keypoints, etc.
+  num_coords: int
+  # The offset of keypoint coordinates in the location tensor.
+  keypoint_coord_offset: int
+  # The number of predicted keypoints.
+  num_keypoints: int
+  # The dimension of each keypoint, e.g. number of values predicted for each
+  # keypoint.
+  num_values_per_keypoint: int
+  # Parameters for decoding SSD detection model.
+  x_scale: float
+  y_scale: float
+  w_scale: float
+  h_scale: float
+  # Whether to apply exponential on box size.
+  apply_exponential_on_box_size: bool
+  # Whether to apply sigmod function on the score.
+  sigmoid_score: bool
+
+
+# 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 ObjectDetectorOptionsMd(metadata_info.CustomMetadataMd):
+  """Object detector options metadata."""
+
+  _METADATA_FILE_IDENTIFIER = b"V001"
+
+  def __init__(
+      self,
+      ssd_anchors_options: SsdAnchorsOptions,
+      tensors_decoding_options: TensorsDecodingOptions,
+  ) -> None:
+    """Creates an ObjectDetectorOptionsMd object.
+
+    Args:
+      ssd_anchors_options: the ssd anchors options associated to the object
+        detector model.
+      tensors_decoding_options: the tensors decoding options used to decode the
+        object detector model output.
+    """
+    if ssd_anchors_options.fixed_anchors_schema is None:
+      raise ValueError(
+          "Currently only support FixedAnchorsSchema, which cannot be found"
+          " in ssd_anchors_options."
+      )
+    self.ssd_anchors_options = ssd_anchors_options
+    self.tensors_decoding_options = tensors_decoding_options
+    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 object
+      detector options metadata.
+    """
+    detector_options = _detector_metadata_fb.ObjectDetectorOptionsT()
+
+    # Set ssd_anchors_options.
+    ssd_anchors_options = _detector_metadata_fb.SsdAnchorsOptionsT()
+    fixed_anchors_schema = _detector_metadata_fb.FixedAnchorsSchemaT()
+    fixed_anchors_schema.anchors = []
+    for anchor in self.ssd_anchors_options.fixed_anchors_schema.anchors:
+      anchor_t = _detector_metadata_fb.FixedAnchorT()
+      anchor_t.xCenter = anchor.x_center
+      anchor_t.yCenter = anchor.y_center
+      anchor_t.width = anchor.width
+      anchor_t.height = anchor.height
+      fixed_anchors_schema.anchors.append(anchor_t)
+    ssd_anchors_options.fixedAnchorsSchema = fixed_anchors_schema
+    detector_options.ssdAnchorsOptions = ssd_anchors_options
+
+    # Set tensors_decoding_options.
+    tensors_decoding_options = _detector_metadata_fb.TensorsDecodingOptionsT()
+    tensors_decoding_options.numClasses = (
+        self.tensors_decoding_options.num_classes
+    )
+    tensors_decoding_options.numBoxes = self.tensors_decoding_options.num_boxes
+    tensors_decoding_options.numCoords = (
+        self.tensors_decoding_options.num_coords
+    )
+    tensors_decoding_options.keypointCoordOffset = (
+        self.tensors_decoding_options.keypoint_coord_offset
+    )
+    tensors_decoding_options.numKeypoints = (
+        self.tensors_decoding_options.num_keypoints
+    )
+    tensors_decoding_options.numValuesPerKeypoint = (
+        self.tensors_decoding_options.num_values_per_keypoint
+    )
+    tensors_decoding_options.xScale = self.tensors_decoding_options.x_scale
+    tensors_decoding_options.yScale = self.tensors_decoding_options.y_scale
+    tensors_decoding_options.wScale = self.tensors_decoding_options.w_scale
+    tensors_decoding_options.hScale = self.tensors_decoding_options.h_scale
+    tensors_decoding_options.applyExponentialOnBoxSize = (
+        self.tensors_decoding_options.apply_exponential_on_box_size
+    )
+    tensors_decoding_options.sigmoidScore = (
+        self.tensors_decoding_options.sigmoid_score
+    )
+    detector_options.tensorsDecodingOptions = tensors_decoding_options
+
+    # Get the object detector options flatbuffer.
+    b = flatbuffers.Builder(0)
+    b.Finish(detector_options.Pack(b), self._METADATA_FILE_IDENTIFIER)
+    detector_options_buf = b.Output()
+
+    # Add the object detector options flatbuffer in custom metadata.
+    custom_metadata = _metadata_fb.CustomMetadataT()
+    custom_metadata.name = self.name
+    custom_metadata.data = detector_options_buf
+    return custom_metadata
+
+
+class MetadataWriter(metadata_writer.MetadataWriterBase):
+  """MetadataWriter to write the metadata into the object detector."""
+
+  @classmethod
+  def create_for_models_with_nms(
+      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 object detector with postprocessing in the model.
+
+    This method create a metadata writer for the models with postprocessing [1].
+
+    The parameters required in this method are mandatory when using MediaPipe
+    Tasks.
+
+    Example usage:
+      metadata_writer = object_detector.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
+        [2].
+      input_norm_std: the std value used in the input tensor normalizarion [2].
+      labels: an instance of Labels helper class used in the output
+        classification tensor [3].
+      score_calibration: A container of the score calibration operation [4] in
+        the classification tensor. Optional if the model does not use score
+        calibration.
+      [1]:
+        https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/detection_postprocess.cc
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
+      [3]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
+      [4]:
+        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_detection_output(labels, score_calibration)
+    return cls(writer)
+
+  @classmethod
+  def create_for_models_without_nms(
+      cls,
+      model_buffer: bytearray,
+      input_norm_mean: List[float],
+      input_norm_std: List[float],
+      labels: metadata_writer.Labels,
+      ssd_anchors_options: SsdAnchorsOptions,
+      tensors_decoding_options: TensorsDecodingOptions,
+      output_tensors_order: metadata_info.RawDetectionOutputTensorsOrder = metadata_info.RawDetectionOutputTensorsOrder.UNSPECIFIED,
+  ) -> "MetadataWriter":
+    """Creates MetadataWriter to write the metadata for object detector without postprocessing in the model.
+
+    This method create a metadata writer for the models without postprocessing
+    [1].
+
+    The parameters required in this method are mandatory when using MediaPipe
+    Tasks.
+
+    Example usage:
+      metadata_writer = object_detector.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
+        [2].
+      input_norm_std: the std value used in the input tensor normalizarion [2].
+      labels: an instance of Labels helper class used in the output
+        classification tensor [3].
+      ssd_anchors_options: the ssd anchors options associated to the object
+        detector model.
+      tensors_decoding_options: the tensors decoding options used to decode the
+        object detector model output.
+      output_tensors_order: the order of the output tensors.
+      [1]:
+        https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/detection_postprocess.cc
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L389
+      [3]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
+
+    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_raw_detection_output(
+        labels, output_tensors_order=output_tensors_order
+    )
+    option_md = ObjectDetectorOptionsMd(
+        ssd_anchors_options, tensors_decoding_options
+    )
+    writer.add_custom_metadata(option_md)
+    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

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

@@ -0,0 +1,119 @@
+# 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 Text classifier models."""
+
+from typing import Union
+
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
+
+_MODEL_NAME = "TextClassifier"
+_MODEL_DESCRIPTION = ("Classify the input text into a set of known categories.")
+
+# The input tensor names of models created by Model Maker.
+_DEFAULT_ID_NAME = "serving_default_input_word_ids:0"
+_DEFAULT_MASK_NAME = "serving_default_input_mask:0"
+_DEFAULT_SEGMENT_ID_NAME = "serving_default_input_type_ids:0"
+
+
+class MetadataWriter(metadata_writer.MetadataWriterBase):
+  """MetadataWriter to write the metadata into the text classifier."""
+
+  @classmethod
+  def create_for_regex_model(
+      cls, model_buffer: bytearray,
+      regex_tokenizer: metadata_writer.RegexTokenizer,
+      labels: metadata_writer.Labels) -> "MetadataWriter":
+    """Creates MetadataWriter for TFLite model with regex tokentizer.
+
+    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.
+      regex_tokenizer: information of the regex tokenizer [1] used to process
+        the input string. If the tokenizer is `BertTokenizer` [2] or
+        `SentencePieceTokenizer` [3], please refer to
+        `create_for_bert_model`.
+      labels: an instance of Labels helper class used in the output
+        classification tensor [4].
+
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
+      [3]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
+      [4]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
+
+    Returns:
+      A MetadataWriter object.
+    """
+    writer = metadata_writer.MetadataWriter(model_buffer)
+    writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
+    writer.add_regex_text_input(regex_tokenizer)
+    writer.add_classification_output(labels)
+    return cls(writer)
+
+  @classmethod
+  def create_for_bert_model(
+      cls,
+      model_buffer: bytearray,
+      tokenizer: Union[metadata_writer.BertTokenizer,
+                       metadata_writer.SentencePieceTokenizer],
+      labels: metadata_writer.Labels,
+      ids_name: str = _DEFAULT_ID_NAME,
+      mask_name: str = _DEFAULT_MASK_NAME,
+      segment_name: str = _DEFAULT_SEGMENT_ID_NAME,
+  ) -> "MetadataWriter":
+    """Creates MetadataWriter for models with {Bert/SentencePiece}Tokenizer.
+
+    `ids_name`, `mask_name`, and `segment_name` correspond to the `Tensor.name`
+    in the TFLite schema, which help to determine the tensor order when
+    populating metadata. The default values come from Model Maker.
+
+    Args:
+      model_buffer: valid buffer of the model file.
+      tokenizer: information of the tokenizer used to process the input string,
+        if any. Supported tokenziers are: `BertTokenizer` [1] and
+        `SentencePieceTokenizer` [2]. If the tokenizer is `RegexTokenizer` [3],
+        refer to `create_for_regex_model`.
+      labels: an instance of Labels helper class used in the output
+        classification tensor [4].
+      ids_name: name of the ids tensor, which represents the tokenized ids of
+        the input text.
+      mask_name: name of the mask tensor, which represents the mask with `1` for
+        real tokens and `0` for padding tokens.
+      segment_name: name of the segment ids tensor, where `0` stands for the
+        first sequence, and `1` stands for the second sequence if exists. [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
+          [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
+          [3]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
+          [4]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L99
+
+    Returns:
+      A MetadataWriter object.
+    """
+    writer = metadata_writer.MetadataWriter(model_buffer)
+    writer.add_general_info(_MODEL_NAME, _MODEL_DESCRIPTION)
+    writer.add_bert_text_input(tokenizer, ids_name, mask_name, segment_name)
+    writer.add_classification_output(labels)
+    return cls(writer)

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

@@ -0,0 +1,91 @@
+# 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.
+# ==============================================================================
+"""Helper methods for writing metadata into TFLite models."""
+
+from typing import List
+
+from mediapipe.tasks.metadata import schema_py_generated as _schema_fb
+
+
+def get_input_tensor_names(model_buffer: bytearray) -> List[str]:
+  """Gets a list of the input tensor names."""
+  subgraph = get_subgraph(model_buffer)
+  tensor_names = []
+  for i in range(subgraph.InputsLength()):
+    index = subgraph.Inputs(i)
+    tensor_names.append(subgraph.Tensors(index).Name().decode("utf-8"))
+  return tensor_names
+
+
+def get_output_tensor_names(model_buffer: bytearray) -> List[str]:
+  """Gets a list of the output tensor names."""
+  subgraph = get_subgraph(model_buffer)
+  tensor_names = []
+  for i in range(subgraph.OutputsLength()):
+    index = subgraph.Outputs(i)
+    tensor_names.append(subgraph.Tensors(index).Name().decode("utf-8"))
+  return tensor_names
+
+
+def get_input_tensor_types(
+    model_buffer: bytearray) -> List[_schema_fb.TensorType]:
+  """Gets a list of the input tensor types."""
+  subgraph = get_subgraph(model_buffer)
+  tensor_types = []
+  for i in range(subgraph.InputsLength()):
+    index = subgraph.Inputs(i)
+    tensor_types.append(subgraph.Tensors(index).Type())
+  return tensor_types
+
+
+def get_output_tensor_types(
+    model_buffer: bytearray) -> List[_schema_fb.TensorType]:
+  """Gets a list of the output tensor types."""
+  subgraph = get_subgraph(model_buffer)
+  tensor_types = []
+  for i in range(subgraph.OutputsLength()):
+    index = subgraph.Outputs(i)
+    tensor_types.append(subgraph.Tensors(index).Type())
+  return tensor_types
+
+
+def get_output_tensor_indices(model_buffer: bytearray) -> List[int]:
+  """Gets a list of the output tensor indices."""
+  subgraph = get_subgraph(model_buffer)
+  return subgraph.OutputsAsNumpy()
+
+
+def get_subgraph(model_buffer: bytearray) -> _schema_fb.SubGraph:
+  """Gets the subgraph of the model.
+
+  TFLite does not support multi-subgraph. A model should have exactly one
+  subgraph.
+
+  Args:
+    model_buffer: valid buffer of the model file.
+
+  Returns:
+    The subgraph of the model.
+
+  Raises:
+    ValueError: if the model has more than one subgraph or has no subgraph.
+  """
+
+  model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
+
+  # Use the first subgraph as default. TFLite Interpreter doesn't support
+  # multiple subgraphs yet, but models with mini-benchmark may have multiple
+  # subgraphs for acceleration evaluation purpose.
+  return model.Subgraphs(0)

mediapipe/tasks/python/test/__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/test/audio/__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.