mediapipe-nightly 0.10.21.post20241223__cp310-cp310-manylinux_2_28_x86_64.whl

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

@@ -0,0 +1,1166 @@
+# 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 classes for common model metadata information."""
+
+import abc
+import collections
+import csv
+import enum
+import os
+from typing import List, Optional, Type, Union
+
+from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
+from mediapipe.tasks.metadata import schema_py_generated as _schema_fb
+from mediapipe.tasks.python.metadata.metadata_writers import writer_utils
+
+# Min and max values for UINT8 tensors.
+_MIN_UINT8 = 0
+_MAX_UINT8 = 255
+
+# Default description for vocabulary files.
+_VOCAB_FILE_DESCRIPTION = ("Vocabulary file to convert natural language "
+                           "words to embedding vectors.")
+
+
+class GeneralMd:
+  """A container for common metadata information of a model.
+
+  Attributes:
+    name: name of the model.
+    version: version of the model.
+    description: description of what the model does.
+    author: author of the model.
+    licenses: licenses of the model.
+  """
+
+  def __init__(self,
+               name: Optional[str] = None,
+               version: Optional[str] = None,
+               description: Optional[str] = None,
+               author: Optional[str] = None,
+               licenses: Optional[str] = None) -> None:
+    self.name = name
+    self.version = version
+    self.description = description
+    self.author = author
+    self.licenses = licenses
+
+  def create_metadata(self) -> _metadata_fb.ModelMetadataT:
+    """Creates the model metadata based on the general model information.
+
+    Returns:
+      A Flatbuffers Python object of the model metadata.
+    """
+    model_metadata = _metadata_fb.ModelMetadataT()
+    model_metadata.name = self.name
+    model_metadata.version = self.version
+    model_metadata.description = self.description
+    model_metadata.author = self.author
+    model_metadata.license = self.licenses
+    return model_metadata
+
+
+class AssociatedFileMd:
+  """A container for common associated file metadata information.
+
+  Attributes:
+    file_path: path to the associated file.
+    description: description of the associated file.
+    file_type: file type of the associated file [1].
+    locale: locale of the associated file [2].
+    [1]:
+      https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L77
+    [2]:
+      https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L176
+  """
+
+  def __init__(
+      self,
+      file_path: str,
+      description: Optional[str] = None,
+      file_type: Optional[int] = _metadata_fb.AssociatedFileType.UNKNOWN,
+      locale: Optional[str] = None) -> None:
+    self.file_path = file_path
+    self.description = description
+    self.file_type = file_type
+    self.locale = locale
+
+  def create_metadata(self) -> _metadata_fb.AssociatedFileT:
+    """Creates the associated file metadata.
+
+    Returns:
+      A Flatbuffers Python object of the associated file metadata.
+    """
+    file_metadata = _metadata_fb.AssociatedFileT()
+    file_metadata.name = os.path.basename(self.file_path)
+    file_metadata.description = self.description
+    file_metadata.type = self.file_type
+    file_metadata.locale = self.locale
+    return file_metadata
+
+
+class LabelFileMd(AssociatedFileMd):
+  """A container for label file metadata information."""
+
+  _LABEL_FILE_DESCRIPTION = ("Labels for categories that the model can "
+                             "recognize.")
+  _FILE_TYPE = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
+
+  def __init__(self, file_path: str, locale: Optional[str] = None) -> None:
+    """Creates a LabelFileMd object.
+
+    Args:
+      file_path: file_path of the label file.
+      locale: locale of the label file [1].
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L176
+    """
+    super().__init__(file_path, self._LABEL_FILE_DESCRIPTION, self._FILE_TYPE,
+                     locale)
+
+
+class ScoreCalibrationMd:
+  """A container for score calibration [1] metadata information.
+
+  [1]:
+    https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L456
+  """
+
+  _SCORE_CALIBRATION_FILE_DESCRIPTION = (
+      "Contains sigmoid-based score calibration parameters. The main purposes "
+      "of score calibration is to make scores across classes comparable, so "
+      "that a common threshold can be used for all output classes.")
+  _FILE_TYPE = _metadata_fb.AssociatedFileType.TENSOR_AXIS_SCORE_CALIBRATION
+
+  def __init__(self,
+               score_transformation_type: _metadata_fb.ScoreTransformationType,
+               default_score: float, file_path: str) -> None:
+    """Creates a ScoreCalibrationMd object.
+
+    Args:
+      score_transformation_type: type of the function used for transforming the
+        uncalibrated score before applying score calibration.
+      default_score: the default calibrated score to apply if the uncalibrated
+        score is below min_score or if no parameters were specified for a given
+        index.
+      file_path: file_path of the score calibration file [1].
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L133
+
+    Raises:
+      ValueError: if the score_calibration file is malformed.
+    """
+    self._score_transformation_type = score_transformation_type
+    self._default_score = default_score
+    self._file_path = file_path
+
+    # Sanity check the score calibration file.
+    with open(self._file_path) as calibration_file:
+      csv_reader = csv.reader(calibration_file, delimiter=",")
+      for row in csv_reader:
+        if row and len(row) != 3 and len(row) != 4:
+          raise ValueError(
+              f"Expected empty lines or 3 or 4 parameters per line in score"
+              f" calibration file, but got {len(row)}.")
+
+        if row and float(row[0]) < 0:
+          raise ValueError(
+              f"Expected scale to be a non-negative value, but got "
+              f"{float(row[0])}.")
+
+  def create_metadata(self) -> _metadata_fb.ProcessUnitT:
+    """Creates the score calibration metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the score calibration metadata.
+    """
+    score_calibration = _metadata_fb.ProcessUnitT()
+    score_calibration.optionsType = (
+        _metadata_fb.ProcessUnitOptions.ScoreCalibrationOptions)
+    options = _metadata_fb.ScoreCalibrationOptionsT()
+    options.scoreTransformation = self._score_transformation_type
+    options.defaultScore = self._default_score
+    score_calibration.options = options
+    return score_calibration
+
+  def create_score_calibration_file_md(self) -> AssociatedFileMd:
+    return AssociatedFileMd(self._file_path,
+                            self._SCORE_CALIBRATION_FILE_DESCRIPTION,
+                            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 RegexTokenizerMd:
+  """A container for the Regex tokenizer [1] metadata information.
+
+  [1]:
+    https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L500
+  """
+
+  def __init__(self, delim_regex_pattern: str, vocab_file_path: str):
+    """Initializes a RegexTokenizerMd object.
+
+    Args:
+      delim_regex_pattern: the regular expression to segment strings and create
+        tokens.
+      vocab_file_path: path to the vocabulary file.
+    """
+    self._delim_regex_pattern = delim_regex_pattern
+    self._vocab_file_path = vocab_file_path
+
+  def create_metadata(self) -> _metadata_fb.ProcessUnitT:
+    """Creates the Regex tokenizer metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the Regex tokenizer metadata.
+    """
+    vocab = _metadata_fb.AssociatedFileT()
+    vocab.name = self._vocab_file_path
+    vocab.description = _VOCAB_FILE_DESCRIPTION
+    vocab.type = _metadata_fb.AssociatedFileType.VOCABULARY
+
+    # Create the RegexTokenizer.
+    tokenizer = _metadata_fb.ProcessUnitT()
+    tokenizer.optionsType = (
+        _metadata_fb.ProcessUnitOptions.RegexTokenizerOptions)
+    tokenizer.options = _metadata_fb.RegexTokenizerOptionsT()
+    tokenizer.options.delimRegexPattern = self._delim_regex_pattern
+    tokenizer.options.vocabFile = [vocab]
+    return tokenizer
+
+
+class BertTokenizerMd:
+  """A container for the Bert tokenizer [1] metadata information.
+
+  [1]:
+    https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
+  """
+
+  def __init__(self, vocab_file_path: str):
+    """Initializes a BertTokenizerMd object.
+
+    Args:
+      vocab_file_path: path to the vocabulary file.
+    """
+    self._vocab_file_path = vocab_file_path
+
+  def create_metadata(self) -> _metadata_fb.ProcessUnitT:
+    """Creates the Bert tokenizer metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the Bert tokenizer metadata.
+    """
+    vocab = _metadata_fb.AssociatedFileT()
+    vocab.name = self._vocab_file_path
+    vocab.description = _VOCAB_FILE_DESCRIPTION
+    vocab.type = _metadata_fb.AssociatedFileType.VOCABULARY
+    tokenizer = _metadata_fb.ProcessUnitT()
+    tokenizer.optionsType = _metadata_fb.ProcessUnitOptions.BertTokenizerOptions
+    tokenizer.options = _metadata_fb.BertTokenizerOptionsT()
+    tokenizer.options.vocabFile = [vocab]
+    return tokenizer
+
+
+class SentencePieceTokenizerMd:
+  """A container for the sentence piece tokenizer [1] metadata information.
+
+  [1]:
+    https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
+  """
+
+  _SP_MODEL_DESCRIPTION = "The sentence piece model file."
+  _SP_VOCAB_FILE_DESCRIPTION = _VOCAB_FILE_DESCRIPTION + (
+      " This file is optional during tokenization, while the sentence piece "
+      "model is mandatory.")
+
+  def __init__(self,
+               sentence_piece_model_path: str,
+               vocab_file_path: Optional[str] = None):
+    """Initializes a SentencePieceTokenizerMd object.
+
+    Args:
+      sentence_piece_model_path: path to the sentence piece model file.
+      vocab_file_path: path to the vocabulary file.
+    """
+    self._sentence_piece_model_path = sentence_piece_model_path
+    self._vocab_file_path = vocab_file_path
+
+  def create_metadata(self) -> _metadata_fb.ProcessUnitT:
+    """Creates the sentence piece tokenizer metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the sentence piece tokenizer metadata.
+    """
+    tokenizer = _metadata_fb.ProcessUnitT()
+    tokenizer.optionsType = (
+        _metadata_fb.ProcessUnitOptions.SentencePieceTokenizerOptions)
+    tokenizer.options = _metadata_fb.SentencePieceTokenizerOptionsT()
+
+    sp_model = _metadata_fb.AssociatedFileT()
+    sp_model.name = self._sentence_piece_model_path
+    sp_model.description = self._SP_MODEL_DESCRIPTION
+    tokenizer.options.sentencePieceModel = [sp_model]
+    if self._vocab_file_path:
+      vocab = _metadata_fb.AssociatedFileT()
+      vocab.name = self._vocab_file_path
+      vocab.description = self._SP_VOCAB_FILE_DESCRIPTION
+      vocab.type = _metadata_fb.AssociatedFileType.VOCABULARY
+      tokenizer.options.vocabFile = [vocab]
+    return tokenizer
+
+
+class ValueRangeMd:
+  """A container for value range metadata information."""
+
+  def __init__(self, min_value: int, max_value: int) -> None:
+    self.min_value = min_value
+    self.max_value = max_value
+
+  def create_metadata(self) -> _metadata_fb.ValueRangeT:
+    """Creates the value range metadata based on the information."""
+    value_range_metadata = _metadata_fb.ValueRangeT()
+    value_range_metadata.min = self.min_value
+    value_range_metadata.max = self.max_value
+    return value_range_metadata
+
+
+class TensorMd:
+  """A container for common tensor metadata information.
+
+  Attributes:
+    name: name of the tensor.
+    description: description of what the tensor is.
+    min_values: per-channel minimum value of the tensor.
+    max_values: per-channel maximum value of the tensor.
+    content_type: content_type of the tensor.
+    associated_files: information of the associated files in the tensor.
+    tensor_name: name of the corresponding tensor [1] in the TFLite model. It is
+      used to locate the corresponding tensor and decide the order of the tensor
+      metadata [2] when populating model metadata.
+    content_range_md: information of content range [3].
+    [1]:
+      https://github.com/tensorflow/tensorflow/blob/cb67fef35567298b40ac166b0581cd8ad68e5a3a/tensorflow/lite/schema/schema.fbs#L1129-L1136
+    [2]:
+      https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L623-L640
+    [3]:
+      https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L385
+  """
+
+  def __init__(
+      self,
+      name: Optional[str] = None,
+      description: Optional[str] = None,
+      min_values: Optional[List[float]] = None,
+      max_values: Optional[List[float]] = None,
+      content_type: int = _metadata_fb.ContentProperties.FeatureProperties,
+      associated_files: Optional[List[Type[AssociatedFileMd]]] = None,
+      tensor_name: Optional[str] = None,
+      content_range_md: Optional[ValueRangeMd] = None,
+  ) -> None:
+    self.name = name
+    self.description = description
+    self.min_values = min_values
+    self.max_values = max_values
+    self.content_type = content_type
+    self.associated_files = associated_files
+    self.tensor_name = tensor_name
+    self.content_range_md = content_range_md
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the input tensor metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the input metadata.
+    """
+    tensor_metadata = _metadata_fb.TensorMetadataT()
+    tensor_metadata.name = self.name
+    tensor_metadata.description = self.description
+
+    # Create min and max values
+    stats = _metadata_fb.StatsT()
+    stats.max = self.max_values
+    stats.min = self.min_values
+    tensor_metadata.stats = stats
+
+    # Create content properties
+    content = _metadata_fb.ContentT()
+    if self.content_type is _metadata_fb.ContentProperties.FeatureProperties:
+      content.contentProperties = _metadata_fb.FeaturePropertiesT()
+    elif self.content_type is _metadata_fb.ContentProperties.ImageProperties:
+      content.contentProperties = _metadata_fb.ImagePropertiesT()
+    elif self.content_type is (
+        _metadata_fb.ContentProperties.BoundingBoxProperties):
+      content.contentProperties = _metadata_fb.BoundingBoxPropertiesT()
+    elif self.content_type is _metadata_fb.ContentProperties.AudioProperties:
+      content.contentProperties = _metadata_fb.AudioPropertiesT()
+
+    content.contentPropertiesType = self.content_type
+    tensor_metadata.content = content
+    if self.content_range_md:
+      tensor_metadata.content.range = self.content_range_md.create_metadata()
+
+    # TODO: check if multiple label files have populated locale.
+    # Create associated files
+    if self.associated_files:
+      tensor_metadata.associatedFiles = [
+          file.create_metadata() for file in self.associated_files
+      ]
+    return tensor_metadata
+
+
+class InputImageTensorMd(TensorMd):
+  """A container for input image tensor metadata information.
+
+  Attributes:
+    norm_mean: the mean value used in tensor normalization [1].
+    norm_std: the std value used in the tensor normalization [1]. norm_mean and
+      norm_std must have the same dimension.
+    color_space_type: the color space type of the input image [2].
+    [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#L198
+  """
+
+  # Min and max float values for image pixels.
+  _MIN_PIXEL = 0.0
+  _MAX_PIXEL = 255.0
+
+  def __init__(
+      self,
+      name: Optional[str] = None,
+      description: Optional[str] = None,
+      norm_mean: Optional[List[float]] = None,
+      norm_std: Optional[List[float]] = None,
+      color_space_type: Optional[int] = _metadata_fb.ColorSpaceType.UNKNOWN,
+      tensor_type: Optional["_schema_fb.TensorType"] = None) -> None:
+    """Initializes the instance of InputImageTensorMd.
+
+    Args:
+      name: name of the tensor.
+      description: description of what the tensor is.
+      norm_mean: the mean value used in tensor normalization [1].
+      norm_std: the std value used in the tensor normalization [1]. norm_mean
+        and norm_std must have the same dimension.
+      color_space_type: the color space type of the input image [2].
+      tensor_type: data type of the tensor.
+      [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#L198
+
+    Raises:
+      ValueError: if norm_mean and norm_std have different dimensions.
+    """
+    if norm_std and norm_mean and len(norm_std) != len(norm_mean):
+      raise ValueError(
+          f"norm_mean and norm_std are expected to be the same dim. But got "
+          f"{len(norm_mean)} and {len(norm_std)}")
+
+    if tensor_type is _schema_fb.TensorType.UINT8:
+      min_values = [_MIN_UINT8]
+      max_values = [_MAX_UINT8]
+    elif tensor_type is _schema_fb.TensorType.FLOAT32 and norm_std and norm_mean:
+      min_values = [
+          float(self._MIN_PIXEL - mean) / std
+          for mean, std in zip(norm_mean, norm_std)
+      ]
+      max_values = [
+          float(self._MAX_PIXEL - mean) / std
+          for mean, std in zip(norm_mean, norm_std)
+      ]
+    else:
+      # Uint8 and Float32 are the two major types currently. And Task library
+      # doesn't support other types so far.
+      min_values = None
+      max_values = None
+
+    super().__init__(name, description, min_values, max_values,
+                     _metadata_fb.ContentProperties.ImageProperties)
+    self.norm_mean = norm_mean
+    self.norm_std = norm_std
+    self.color_space_type = color_space_type
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the input image metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the input image metadata.
+    """
+    tensor_metadata = super().create_metadata()
+    tensor_metadata.content.contentProperties.colorSpace = self.color_space_type
+    # Create normalization parameters
+    if self.norm_mean and self.norm_std:
+      normalization = _metadata_fb.ProcessUnitT()
+      normalization.optionsType = (
+          _metadata_fb.ProcessUnitOptions.NormalizationOptions)
+      normalization.options = _metadata_fb.NormalizationOptionsT()
+      normalization.options.mean = self.norm_mean
+      normalization.options.std = self.norm_std
+      tensor_metadata.processUnits = [normalization]
+    return tensor_metadata
+
+
+class InputTextTensorMd(TensorMd):
+  """A container for the input text tensor metadata information.
+
+  Attributes:
+    tokenizer_md: information of the tokenizer in the input text tensor, if any.
+  """
+
+  def __init__(self,
+               name: Optional[str] = None,
+               description: Optional[str] = None,
+               tokenizer_md: Optional[RegexTokenizerMd] = None):
+    """Initializes the instance of InputTextTensorMd.
+
+    Args:
+      name: name of the tensor.
+      description: description of what the tensor is.
+      tokenizer_md: information of the tokenizer in the input text tensor, if
+        any. Only `RegexTokenizer` [1] is currently supported. If the tokenizer
+        is `BertTokenizer` [2] or `SentencePieceTokenizer` [3], refer to
+        `BertInputTensorsMd` class.
+      [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
+    """
+    super().__init__(name, description)
+    self.tokenizer_md = tokenizer_md
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the input text metadata based on the information.
+
+    Returns:
+      A Flatbuffers Python object of the input text metadata.
+
+    Raises:
+      ValueError: if the type of tokenizer_md is unsupported.
+    """
+    if not isinstance(self.tokenizer_md, (type(None), RegexTokenizerMd)):
+      raise ValueError(
+          f"The type of tokenizer_options, {type(self.tokenizer_md)}, is "
+          f"unsupported")
+
+    tensor_metadata = super().create_metadata()
+    if self.tokenizer_md:
+      tensor_metadata.processUnits = [self.tokenizer_md.create_metadata()]
+    return tensor_metadata
+
+
+def _get_file_paths(files: List[_metadata_fb.AssociatedFileT]) -> List[str]:
+  """Gets file paths from a list of associated files."""
+  if not files:
+    return []
+  return [file.name for file in files]
+
+
+def _get_tokenizer_associated_files(
+    tokenizer_options: Optional[
+        Union[_metadata_fb.BertTokenizerOptionsT,
+              _metadata_fb.SentencePieceTokenizerOptionsT]]
+) -> List[str]:
+  """Gets a list of associated files packed in the tokenizer_options.
+
+  Args:
+    tokenizer_options: a tokenizer metadata object. Support the following
+      tokenizer types:
+      1. BertTokenizerOptions:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L477
+      2. SentencePieceTokenizerOptions:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L485
+
+  Returns:
+    A list of associated files included in tokenizer_options.
+  """
+
+  if not tokenizer_options:
+    return []
+
+  if isinstance(tokenizer_options, _metadata_fb.BertTokenizerOptionsT):
+    return _get_file_paths(tokenizer_options.vocabFile)
+  elif isinstance(tokenizer_options,
+                  _metadata_fb.SentencePieceTokenizerOptionsT):
+    return _get_file_paths(tokenizer_options.vocabFile) + _get_file_paths(
+        tokenizer_options.sentencePieceModel)
+  else:
+    return []
+
+
+class BertInputTensorsMd:
+  """A container for the input tensor metadata information of Bert models."""
+
+  _IDS_NAME = "ids"
+  _IDS_DESCRIPTION = "Tokenized ids of the input text."
+  _MASK_NAME = "mask"
+  _MASK_DESCRIPTION = ("Mask with 1 for real tokens and 0 for padding "
+                       "tokens.")
+  _SEGMENT_IDS_NAME = "segment_ids"
+  _SEGMENT_IDS_DESCRIPTION = (
+      "0 for the first sequence, 1 for the second sequence if exists.")
+
+  def __init__(self,
+               model_buffer: bytearray,
+               ids_name: str,
+               mask_name: str,
+               segment_name: str,
+               tokenizer_md: Union[None, BertTokenizerMd,
+                                   SentencePieceTokenizerMd] = None):
+    """Initializes a BertInputTensorsMd object.
+
+    `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.
+
+    Args:
+      model_buffer: valid buffer of the model file.
+      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.
+      tokenizer_md: information of the tokenizer used to process the input
+        string, if any. Supported tokenizers are: `BertTokenizer` [1] and
+        `SentencePieceTokenizer` [2]. If the tokenizer is `RegexTokenizer` [3],
+        refer to `InputTensorsMd`.
+      [1]:
+        https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L436
+      [2]:
+        https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L473
+      [3]:
+        https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L475
+    """
+    # Verify that tflite_input_names (read from the model) and
+    # input_name (collected from users) are aligned.
+    tflite_input_names = writer_utils.get_input_tensor_names(model_buffer)
+    input_names = [ids_name, mask_name, segment_name]
+    if collections.Counter(tflite_input_names) != collections.Counter(
+        input_names):
+      raise ValueError(
+          f"The input tensor names ({input_names}) do not match the tensor "
+          f"names read from the model ({tflite_input_names}).")
+
+    ids_md = TensorMd(
+        name=self._IDS_NAME,
+        description=self._IDS_DESCRIPTION,
+        tensor_name=ids_name)
+
+    mask_md = TensorMd(
+        name=self._MASK_NAME,
+        description=self._MASK_DESCRIPTION,
+        tensor_name=mask_name)
+
+    segment_ids_md = TensorMd(
+        name=self._SEGMENT_IDS_NAME,
+        description=self._SEGMENT_IDS_DESCRIPTION,
+        tensor_name=segment_name)
+
+    self._input_md = [ids_md, mask_md, segment_ids_md]
+
+    if not isinstance(tokenizer_md,
+                      (type(None), BertTokenizerMd, SentencePieceTokenizerMd)):
+      raise ValueError(
+          f"The type of tokenizer_options, {type(tokenizer_md)}, is unsupported"
+      )
+
+    self._tokenizer_md = tokenizer_md
+
+  def create_input_process_unit_metadata(
+      self) -> List[_metadata_fb.ProcessUnitT]:
+    """Creates the input process unit metadata."""
+    if self._tokenizer_md:
+      return [self._tokenizer_md.create_metadata()]
+    else:
+      return []
+
+  def get_tokenizer_associated_files(self) -> List[str]:
+    """Gets the associated files that are packed in the tokenizer."""
+    if self._tokenizer_md:
+      return _get_tokenizer_associated_files(
+          self._tokenizer_md.create_metadata().options)
+    else:
+      return []
+
+  @property
+  def input_md(self) -> List[TensorMd]:
+    return self._input_md
+
+
+class ClassificationTensorMd(TensorMd):
+  """A container for the classification tensor metadata information.
+
+  Attributes:
+    label_files: information of the label files [1] in the classification
+      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,
+      score_thresholding_md: Optional[ScoreThresholdingMd] = None,
+      content_range_md: Optional[ValueRangeMd] = None,
+  ) -> None:
+    """Initializes the instance of ClassificationTensorMd.
+
+    Args:
+      name: name of the tensor.
+      description: description of what the tensor is.
+      label_files: information of the label files [1] in the classification
+        tensor.
+      tensor_type: data type of the tensor.
+      score_calibration_md: information of the score calibration files operation
+        [2] in the classification tensor.
+      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.
+      content_range_md: information of content range [6].
+      [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/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
+      [6]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L385
+    """
+    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]
+      max_values = [_MAX_UINT8]
+    elif tensor_type is _schema_fb.TensorType.FLOAT32:
+      min_values = [self._MIN_FLOAT]
+      max_values = [self._MAX_FLOAT]
+    else:
+      # Uint8 and Float32 are the two major types currently. And Task library
+      # doesn't support other types so far.
+      min_values = None
+      max_values = None
+
+    associated_files = label_files or []
+    if self.score_calibration_md:
+      associated_files.append(
+          score_calibration_md.create_score_calibration_file_md())
+
+    super().__init__(
+        name,
+        description,
+        min_values,
+        max_values,
+        _metadata_fb.ContentProperties.FeatureProperties,
+        associated_files,
+        tensor_name,
+        content_range_md,
+    )
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the classification tensor metadata based on the information."""
+    tensor_metadata = super().create_metadata()
+    if self.score_calibration_md:
+      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
+
+
+class LocationTensorMd(TensorMd):
+  """A container for the detection location tensor metadata information."""
+
+  # The default order is {left, top, right, bottom}. Denote the order to be
+  # {top, left, bottom, right}.
+  _BOUNDING_BOX_INDEX = (1, 0, 3, 2)
+
+  def __init__(
+      self,
+      name: Optional[str] = None,
+      description: Optional[str] = None,
+      content_range_md: Optional[ValueRangeMd] = None,
+  ) -> None:
+    super().__init__(
+        name=name, description=description, content_range_md=content_range_md
+    )
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the detection location tensor metadata."""
+    content = _metadata_fb.ContentT()
+    content.contentPropertiesType = (
+        _metadata_fb.ContentProperties.BoundingBoxProperties
+    )
+    properties = _metadata_fb.BoundingBoxPropertiesT()
+    properties.index = list(self._BOUNDING_BOX_INDEX)
+    properties.type = _metadata_fb.BoundingBoxType.BOUNDARIES
+    properties.coordinateType = _metadata_fb.CoordinateType.RATIO
+    content.contentProperties = properties
+    if self.content_range_md:
+      content.range = self.content_range_md.create_metadata()
+    location_metadata = super().create_metadata()
+    location_metadata.content = content
+    return location_metadata
+
+
+class CategoryTensorMd(TensorMd):
+  """A container for the category tensor metadata information."""
+
+  def __init__(
+      self,
+      name: Optional[str] = None,
+      description: Optional[str] = None,
+      label_files: Optional[List[LabelFileMd]] = None,
+      content_range_md: Optional[ValueRangeMd] = None,
+  ):
+    """Initializes a CategoryTensorMd object.
+
+    Args:
+      name: name of the tensor.
+      description: description of what the tensor is.
+      label_files: information of the label files [1] in the category tensor.
+      content_range_md: information of content range [2].
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L116
+      [2]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L385
+    """
+    # In category tensors, label files are in the type of TENSOR_VALUE_LABELS.
+    if label_files:
+      for file in label_files:
+        file.file_type = _metadata_fb.AssociatedFileType.TENSOR_VALUE_LABELS
+
+    super().__init__(
+        name=name,
+        description=description,
+        associated_files=label_files,
+        content_range_md=content_range_md,
+    )
+
+
+class DetectionOutputTensorsMd:
+  """A container for the output tensor metadata of detection models."""
+
+  _LOCATION_NAME = "location"
+  _LOCATION_DESCRIPTION = "The locations of the detected boxes."
+  _CATRGORY_NAME = "category"
+  _CATEGORY_DESCRIPTION = "The categories of the detected boxes."
+  _SCORE_NAME = "score"
+  _SCORE_DESCRIPTION = "The scores of the detected boxes."
+  _NUMBER_NAME = "number of detections"
+  _NUMBER_DESCRIPTION = "The number of the detected boxes."
+  _CONTENT_VALUE_DIM = 2
+
+  def __init__(
+      self,
+      model_buffer: bytearray,
+      label_files: Optional[List[LabelFileMd]] = None,
+      score_calibration_md: Optional[ScoreCalibrationMd] = None,
+  ) -> None:
+    """Initializes the instance of DetectionOutputTensorsMd.
+
+    Args:
+      model_buffer: A valid flatbuffer loaded from the TFLite model file.
+      label_files: information of the label files [1] in the classification
+        tensor.
+      score_calibration_md: information of the score calibration files operation
+        [2] 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
+    """
+    content_range_md = ValueRangeMd(
+        min_value=self._CONTENT_VALUE_DIM, max_value=self._CONTENT_VALUE_DIM
+    )
+    location_md = LocationTensorMd(
+        name=self._LOCATION_NAME,
+        description=self._LOCATION_DESCRIPTION,
+        content_range_md=content_range_md,
+    )
+    category_md = CategoryTensorMd(
+        name=self._CATRGORY_NAME,
+        description=self._CATEGORY_DESCRIPTION,
+        label_files=label_files,
+        content_range_md=content_range_md,
+    )
+    score_md = ClassificationTensorMd(
+        name=self._SCORE_NAME,
+        description=self._SCORE_DESCRIPTION,
+        score_calibration_md=score_calibration_md,
+        content_range_md=content_range_md,
+    )
+    number_md = TensorMd(
+        name=self._NUMBER_NAME, description=self._NUMBER_DESCRIPTION
+    )
+
+    # Get the tensor indices of tflite outputs and then gets the order of the
+    # output metadata by the value of tensor indices. The output tensor indices
+    # follow the order as [location, category, score,# detections]. For
+    # instance, if the output indices are [601, 599, 598, 600], tensor names and
+    # indices aligned as below:
+    #   - (598, location)
+    #   - (599, category)
+    #   - (600, score)
+    #   - (601, number of detections)
+    # because of the op's ports of TFLITE_DETECTION_POST_PROCESS
+    # (https://github.com/tensorflow/tensorflow/blob/a4fe268ea084e7d323133ed7b986e0ae259a2bc7/tensorflow/lite/kernels/detection_postprocess.cc#L47-L50).
+    # Thus, the metadata of tensors are paired with output tensor indices & name
+    # in this way.
+
+    # Get the output tensor indices and names from the tflite model.
+    tensor_indices_and_names = zip(
+        writer_utils.get_output_tensor_indices(model_buffer),
+        writer_utils.get_output_tensor_names(model_buffer),
+    )
+    # Sort by the output tensor indices.
+    tensor_indices_and_names = sorted(tensor_indices_and_names)
+
+    # Align tensor names with metadata.
+    self._output_mds = [location_md, category_md, score_md, number_md]
+    if len(self._output_mds) != len(tensor_indices_and_names):
+      raise ValueError(
+          "The size of TFLite output should be " + str(len(self._output_mds))
+      )
+    for i, output_md in enumerate(self._output_mds):
+      output_md.tensor_name = tensor_indices_and_names[i][1]
+
+  @property
+  def output_mds(self) -> List[TensorMd]:
+    return self._output_mds
+
+
+class RawDetectionOutputTensorsOrder(enum.Enum):
+  """Output tensors order for detection models without postprocessing.
+
+  Because it is not able to determined the order of output tensors for models
+  without postprocessing, it is needed to specify the output tensors order for
+  metadata writer.
+  """
+
+  UNSPECIFIED = 0
+  # The first tensor is score, and the second tensor is location.
+  SCORE_LOCATION = 1
+  # The first tensor is location, and the second tensor is score.
+  LOCATION_SCORE = 2
+
+
+class RawDetectionOutputTensorsMd:
+  """A container for the output tensor metadata of detection models without postprocessing."""
+
+  _LOCATION_NAME = "location"
+  _LOCATION_DESCRIPTION = "The locations of the detected boxes."
+  _SCORE_NAME = "score"
+  _SCORE_DESCRIPTION = "The scores of the detected boxes."
+  _CONTENT_VALUE_DIM = 2
+
+  def __init__(
+      self,
+      model_buffer: bytearray,
+      label_files: Optional[List[LabelFileMd]] = None,
+      output_tensors_order: RawDetectionOutputTensorsOrder = RawDetectionOutputTensorsOrder.UNSPECIFIED,
+  ) -> None:
+    """Initializes the instance of DetectionOutputTensorsMd.
+
+    Args:
+      model_buffer: A valid flatbuffer loaded from the TFLite model file.
+      label_files: information of the label files [1] in the classification
+        tensor.
+      output_tensors_order: the order of the output tensors.
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L9
+    """
+    # Get the output tensor indices and names from the tflite model.
+    tensor_indices_and_names = list(
+        zip(
+            writer_utils.get_output_tensor_indices(model_buffer),
+            writer_utils.get_output_tensor_names(model_buffer),
+        )
+    )
+    location_md = LocationTensorMd(
+        name=self._LOCATION_NAME,
+        description=self._LOCATION_DESCRIPTION,
+    )
+    score_md = ClassificationTensorMd(
+        name=self._SCORE_NAME,
+        description=self._SCORE_DESCRIPTION,
+        label_files=label_files,
+    )
+
+    if output_tensors_order == RawDetectionOutputTensorsOrder.SCORE_LOCATION:
+      self._output_mds = [score_md, location_md]
+    elif output_tensors_order == RawDetectionOutputTensorsOrder.LOCATION_SCORE:
+      self._output_mds = [location_md, score_md]
+    else:
+      raise ValueError(
+          f"Unsupported OutputTensorsOrder value: {output_tensors_order}"
+      )
+
+    if len(self._output_mds) != len(tensor_indices_and_names):
+      raise ValueError(
+          "The size of TFLite output should be " + str(len(self._output_mds))
+      )
+    for i, output_md in enumerate(self._output_mds):
+      output_md.tensor_name = tensor_indices_and_names[i][1]
+
+  @property
+  def output_mds(self) -> List[TensorMd]:
+    return self._output_mds
+
+
+class TensorGroupMd:
+  """A container for a group of tensor metadata information."""
+
+  def __init__(
+      self, name: Optional[str] = None, tensor_names: Optional[List[str]] = None
+  ) -> None:
+    """Initializes a CategoryTensorMd object.
+
+    Args:
+      name: name of tensor group.
+      tensor_names:  Names of the tensors to group together, corresponding to
+        TensorMetadata.name [1].
+      [1]:
+        https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L564
+    """
+    self.name = name
+    self.tensor_names = tensor_names
+
+  def create_metadata(self) -> _metadata_fb.TensorGroupT:
+    """Creates the tensor group metadata."""
+    group = _metadata_fb.TensorGroupT()
+    group.name = self.name
+    group.tensorNames = self.tensor_names
+    return group
+
+
+class SegmentationMaskMd(TensorMd):
+  """A container for the segmentation mask metadata information."""
+
+  # The output tensor is in the shape of [1, ImageHeight, ImageWidth, N], where
+  # N is the number of objects that the segmentation model can recognize. The
+  # output tensor is essentially a list of grayscale bitmaps, where each value
+  # is the probability of the corresponding pixel belonging to a certain object
+  # type. Therefore, the content dimension range of the output tensor is [1, 2].
+  _CONTENT_DIM_MIN = 1
+  _CONTENT_DIM_MAX = 2
+
+  def __init__(
+      self,
+      name: Optional[str] = None,
+      description: Optional[str] = None,
+      label_files: Optional[List[LabelFileMd]] = None,
+  ):
+    self.name = name
+    self.description = description
+    associated_files = label_files or []
+    super().__init__(
+        name=name, description=description, associated_files=associated_files
+    )
+
+  def create_metadata(self) -> _metadata_fb.TensorMetadataT:
+    """Creates the metadata for the segmentation masks tensor."""
+    masks_metadata = super().create_metadata()
+
+    # Create tensor content information.
+    content = _metadata_fb.ContentT()
+    content.contentProperties = _metadata_fb.ImagePropertiesT()
+    content.contentProperties.colorSpace = _metadata_fb.ColorSpaceType.GRAYSCALE
+    content.contentPropertiesType = (
+        _metadata_fb.ContentProperties.ImageProperties
+    )
+    # Add the content range. See
+    # https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L323-L385
+    dim_range = _metadata_fb.ValueRangeT()
+    dim_range.min = self._CONTENT_DIM_MIN
+    dim_range.max = self._CONTENT_DIM_MAX
+    content.range = dim_range
+    masks_metadata.content = content
+
+    return masks_metadata
+
+
+class CustomMetadataMd(abc.ABC):
+  """An abstract class of a container for the custom metadata information."""
+
+  def __init__(self, name: Optional[str] = None):
+    self.name = name
+
+  @abc.abstractmethod
+  def create_metadata(self) -> _metadata_fb.CustomMetadataT:
+    """Creates the custom metadata based on the information."""