mediapipe-nightly 0.10.21.post20250114__cp39-cp39-manylinux_2_28_x86_64.whl

mediapipe/tasks/python/metadata/metadata.py

@@ -0,0 +1,928 @@
+# 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.
+# ==============================================================================
+"""TensorFlow Lite metadata tools."""
+
+import copy
+import inspect
+import io
+import json
+import logging
+import os
+import shutil
+import sys
+import tempfile
+from typing import Dict, Optional
+import warnings
+import zipfile
+
+import flatbuffers
+
+from mediapipe.tasks.cc.metadata.python import _pywrap_metadata_version
+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.flatbuffers_lib import _pywrap_flatbuffers
+
+try:
+  # If exists, optionally use TensorFlow to open and check files. Used to
+  # support more than local file systems.
+  # In pip requirements, we doesn't necessarily need tensorflow as a dep.
+  import tensorflow as tf
+  _open_file = tf.io.gfile.GFile
+  _exists_file = tf.io.gfile.exists
+except ImportError as e:
+  # If TensorFlow package doesn't exist, fall back to original open and exists.
+  _open_file = open
+  _exists_file = os.path.exists
+
+
+def _maybe_open_as_binary(filename, mode):
+  """Maybe open the binary file, and returns a file-like."""
+  if hasattr(filename, "read"):  # A file-like has read().
+    return filename
+  openmode = mode if "b" in mode else mode + "b"  # Add binary explicitly.
+  return _open_file(filename, openmode)
+
+
+def _open_as_zipfile(filename, mode="r"):
+  """Open file as a zipfile.
+
+  Args:
+    filename: str or file-like or path-like, to the zipfile.
+    mode: str, common file mode for zip.
+          (See: https://docs.python.org/3/library/zipfile.html)
+
+  Returns:
+    A ZipFile object.
+  """
+  file_like = _maybe_open_as_binary(filename, mode)
+  return zipfile.ZipFile(file_like, mode)
+
+
+def _is_zipfile(filename):
+  """Checks whether it is a zipfile."""
+  with _maybe_open_as_binary(filename, "r") as f:
+    return zipfile.is_zipfile(f)
+
+
+def get_path_to_datafile(path):
+  """Gets the path to the specified file in the data dependencies.
+
+  The path is relative to the file calling the function.
+
+  It's a simple replacement of
+  "tensorflow.python.platform.resource_loader.get_path_to_datafile".
+
+  Args:
+    path: a string resource path relative to the calling file.
+
+  Returns:
+    The path to the specified file present in the data attribute of py_test
+    or py_binary.
+  """
+  data_files_path = os.path.dirname(inspect.getfile(sys._getframe(1)))  # pylint: disable=protected-access
+  return os.path.join(data_files_path, path)
+
+
+_FLATC_TFLITE_METADATA_SCHEMA_FILE = get_path_to_datafile(
+    "../../metadata/metadata_schema.fbs")
+
+
+# TODO: add delete method for associated files.
+class MetadataPopulator(object):
+  """Packs metadata and associated files into TensorFlow Lite model file.
+
+  MetadataPopulator can be used to populate metadata and model associated files
+  into a model file or a model buffer (in bytearray). It can also help to
+  inspect list of files that have been packed into the model or are supposed to
+  be packed into the model.
+
+  The metadata file (or buffer) should be generated based on the metadata
+  schema:
+  mediapipe/tasks/metadata/metadata_schema.fbs
+
+  Example usage:
+  Populate metadata and label file into an image classifier model.
+
+  First, based on metadata_schema.fbs, generate the metadata for this image
+  classifier model using Flatbuffers API. Attach the label file onto the output
+  tensor (the tensor of probabilities) in the metadata.
+
+  Then, pack the metadata and label file into the model as follows.
+
+    ```python
+    # Populating a metadata file (or a metadata buffer) and associated files to
+    a model file:
+    populator = MetadataPopulator.with_model_file(model_file)
+    # For metadata buffer (bytearray read from the metadata file), use:
+    # populator.load_metadata_buffer(metadata_buf)
+    populator.load_metadata_file(metadata_file)
+    populator.load_associated_files([label.txt])
+    # For associated file buffer (bytearray read from the file), use:
+    # populator.load_associated_file_buffers({"label.txt": b"file content"})
+    populator.populate()
+
+    # Populating a metadata file (or a metadata buffer) and associated files to
+    a model buffer:
+    populator = MetadataPopulator.with_model_buffer(model_buf)
+    populator.load_metadata_file(metadata_file)
+    populator.load_associated_files([label.txt])
+    populator.populate()
+    # Writing the updated model buffer into a file.
+    updated_model_buf = populator.get_model_buffer()
+    with open("updated_model.tflite", "wb") as f:
+      f.write(updated_model_buf)
+
+    # Transferring metadata and associated files from another TFLite model:
+    populator = MetadataPopulator.with_model_buffer(model_buf)
+    populator_dst.load_metadata_and_associated_files(src_model_buf)
+    populator_dst.populate()
+    updated_model_buf = populator.get_model_buffer()
+    with open("updated_model.tflite", "wb") as f:
+      f.write(updated_model_buf)
+    ```
+
+  Note that existing metadata buffer (if applied) will be overridden by the new
+  metadata buffer.
+  """
+  # As Zip API is used to concatenate associated files after tflite model file,
+  # the populating operation is developed based on a model file. For in-memory
+  # model buffer, we create a tempfile to serve the populating operation.
+  # Creating the deleting such a tempfile is handled by the class,
+  # _MetadataPopulatorWithBuffer.
+
+  METADATA_FIELD_NAME = "TFLITE_METADATA"
+  TFLITE_FILE_IDENTIFIER = b"TFL3"
+  METADATA_FILE_IDENTIFIER = b"M001"
+
+  def __init__(self, model_file):
+    """Constructor for MetadataPopulator.
+
+    Args:
+      model_file: valid path to a TensorFlow Lite model file.
+
+    Raises:
+      IOError: File not found.
+      ValueError: the model does not have the expected flatbuffer identifier.
+    """
+    _assert_model_file_identifier(model_file)
+    self._model_file = model_file
+    self._metadata_buf = None
+    # _associated_files is a dict of file name and file buffer.
+    self._associated_files = {}
+
+  @classmethod
+  def with_model_file(cls, model_file):
+    """Creates a MetadataPopulator object that populates data to a model file.
+
+    Args:
+      model_file: valid path to a TensorFlow Lite model file.
+
+    Returns:
+      MetadataPopulator object.
+
+    Raises:
+      IOError: File not found.
+      ValueError: the model does not have the expected flatbuffer identifier.
+    """
+    return cls(model_file)
+
+  # TODO: investigate if type check can be applied to model_buf for
+  # FB.
+  @classmethod
+  def with_model_buffer(cls, model_buf):
+    """Creates a MetadataPopulator object that populates data to a model buffer.
+
+    Args:
+      model_buf: TensorFlow Lite model buffer in bytearray.
+
+    Returns:
+      A MetadataPopulator(_MetadataPopulatorWithBuffer) object.
+
+    Raises:
+      ValueError: the model does not have the expected flatbuffer identifier.
+    """
+    return _MetadataPopulatorWithBuffer(model_buf)
+
+  def get_model_buffer(self):
+    """Gets the buffer of the model with packed metadata and associated files.
+
+    Returns:
+      Model buffer (in bytearray).
+    """
+    with _open_file(self._model_file, "rb") as f:
+      return f.read()
+
+  def get_packed_associated_file_list(self):
+    """Gets a list of associated files packed to the model file.
+
+    Returns:
+      List of packed associated files.
+    """
+    if not _is_zipfile(self._model_file):
+      return []
+
+    with _open_as_zipfile(self._model_file, "r") as zf:
+      return zf.namelist()
+
+  def get_recorded_associated_file_list(self):
+    """Gets a list of associated files recorded in metadata of the model file.
+
+    Associated files may be attached to a model, a subgraph, or an input/output
+    tensor.
+
+    Returns:
+      List of recorded associated files.
+    """
+    if not self._metadata_buf:
+      return []
+
+    metadata = _metadata_fb.ModelMetadataT.InitFromObj(
+        _metadata_fb.ModelMetadata.GetRootAsModelMetadata(
+            self._metadata_buf, 0))
+
+    return [
+        file.name.decode("utf-8")
+        for file in self._get_recorded_associated_file_object_list(metadata)
+    ]
+
+  def load_associated_file_buffers(self, associated_files):
+    """Loads the associated file buffers (in bytearray) to be populated.
+
+    Args:
+      associated_files: a dictionary of associated file names and corresponding
+        file buffers, such as {"file.txt": b"file content"}. If pass in file
+          paths for the file name, only the basename will be populated.
+    """
+
+    self._associated_files.update({
+        os.path.basename(name): buffers
+        for name, buffers in associated_files.items()
+    })
+
+  def load_associated_files(self, associated_files):
+    """Loads associated files that to be concatenated after the model file.
+
+    Args:
+      associated_files: list of file paths.
+
+    Raises:
+      IOError:
+        File not found.
+    """
+    for af_name in associated_files:
+      _assert_file_exist(af_name)
+      with _open_file(af_name, "rb") as af:
+        self.load_associated_file_buffers({af_name: af.read()})
+
+  def load_metadata_buffer(self, metadata_buf):
+    """Loads the metadata buffer (in bytearray) to be populated.
+
+    Args:
+      metadata_buf: metadata buffer (in bytearray) to be populated.
+
+    Raises:
+      ValueError: The metadata to be populated is empty.
+      ValueError: The metadata does not have the expected flatbuffer identifier.
+      ValueError: Cannot get minimum metadata parser version.
+      ValueError: The number of SubgraphMetadata is not 1.
+      ValueError: The number of input/output tensors does not match the number
+        of input/output tensor metadata.
+    """
+    if not metadata_buf:
+      raise ValueError("The metadata to be populated is empty.")
+
+    self._validate_metadata(metadata_buf)
+
+    # Gets the minimum metadata parser version of the metadata_buf.
+    min_version = _pywrap_metadata_version.GetMinimumMetadataParserVersion(
+        bytes(metadata_buf))
+
+    # Inserts in the minimum metadata parser version into the metadata_buf.
+    metadata = _metadata_fb.ModelMetadataT.InitFromObj(
+        _metadata_fb.ModelMetadata.GetRootAsModelMetadata(metadata_buf, 0))
+    metadata.minParserVersion = min_version
+
+    # Remove local file directory in the `name` field of `AssociatedFileT`, and
+    # make it consistent with the name of the actual file packed in the model.
+    self._use_basename_for_associated_files_in_metadata(metadata)
+
+    b = flatbuffers.Builder(0)
+    b.Finish(metadata.Pack(b), self.METADATA_FILE_IDENTIFIER)
+    metadata_buf_with_version = b.Output()
+
+    self._metadata_buf = metadata_buf_with_version
+
+  def load_metadata_file(self, metadata_file):
+    """Loads the metadata file to be populated.
+
+    Args:
+      metadata_file: path to the metadata file to be populated.
+
+    Raises:
+      IOError: File not found.
+      ValueError: The metadata to be populated is empty.
+      ValueError: The metadata does not have the expected flatbuffer identifier.
+      ValueError: Cannot get minimum metadata parser version.
+      ValueError: The number of SubgraphMetadata is not 1.
+      ValueError: The number of input/output tensors does not match the number
+        of input/output tensor metadata.
+    """
+    _assert_file_exist(metadata_file)
+    with _open_file(metadata_file, "rb") as f:
+      metadata_buf = f.read()
+    self.load_metadata_buffer(bytearray(metadata_buf))
+
+  def load_metadata_and_associated_files(self, src_model_buf):
+    """Loads the metadata and associated files from another model buffer.
+
+    Args:
+      src_model_buf: source model buffer (in bytearray) with metadata and
+        associated files.
+    """
+    # Load the model metadata from src_model_buf if exist.
+    metadata_buffer = get_metadata_buffer(src_model_buf)
+    if metadata_buffer:
+      self.load_metadata_buffer(metadata_buffer)
+
+    # Load the associated files from src_model_buf if exist.
+    if _is_zipfile(io.BytesIO(src_model_buf)):
+      with _open_as_zipfile(io.BytesIO(src_model_buf)) as zf:
+        self.load_associated_file_buffers(
+            {f: zf.read(f) for f in zf.namelist()})
+
+  def populate(self):
+    """Populates loaded metadata and associated files into the model file."""
+    self._assert_validate()
+    self._populate_metadata_buffer()
+    self._populate_associated_files()
+
+  def _assert_validate(self):
+    """Validates the metadata and associated files to be populated.
+
+    Raises:
+      ValueError:
+        File is recorded in the metadata, but is not going to be populated.
+        File has already been packed.
+    """
+    # Gets files that are recorded in metadata.
+    recorded_files = self.get_recorded_associated_file_list()
+
+    # Gets files that have been packed to self._model_file.
+    packed_files = self.get_packed_associated_file_list()
+
+    # Gets the file name of those associated files to be populated.
+    to_be_populated_files = self._associated_files.keys()
+
+    # Checks all files recorded in the metadata will be populated.
+    for rf in recorded_files:
+      if rf not in to_be_populated_files and rf not in packed_files:
+        raise ValueError("File, '{0}', is recorded in the metadata, but has "
+                         "not been loaded into the populator.".format(rf))
+
+    for f in to_be_populated_files:
+      if f in packed_files:
+        raise ValueError("File, '{0}', has already been packed.".format(f))
+
+      if f not in recorded_files:
+        warnings.warn(
+            "File, '{0}', does not exist in the metadata. But packing it to "
+            "tflite model is still allowed.".format(f))
+
+  def _copy_archived_files(self, src_zip, file_list, dst_zip):
+    """Copy archieved files in file_list from src_zip ro dst_zip."""
+
+    if not _is_zipfile(src_zip):
+      raise ValueError("File, '{0}', is not a zipfile.".format(src_zip))
+
+    with _open_as_zipfile(src_zip, "r") as src_zf, \
+         _open_as_zipfile(dst_zip, "a") as dst_zf:
+      src_list = src_zf.namelist()
+      for f in file_list:
+        if f not in src_list:
+          raise ValueError(
+              "File, '{0}', does not exist in the zipfile, {1}.".format(
+                  f, src_zip))
+        file_buffer = src_zf.read(f)
+        dst_zf.writestr(f, file_buffer)
+
+  def _get_associated_files_from_process_units(self, table, field_name):
+    """Gets the files that are attached the process units field of a table.
+
+    Args:
+      table: a Flatbuffers table object that contains fields of an array of
+        ProcessUnit, such as TensorMetadata and SubGraphMetadata.
+      field_name: the name of the field in the table that represents an array of
+        ProcessUnit. If the table is TensorMetadata, field_name can be
+        "ProcessUnits". If the table is SubGraphMetadata, field_name can be
+        either "InputProcessUnits" or "OutputProcessUnits".
+
+    Returns:
+      A list of AssociatedFileT objects.
+    """
+
+    if table is None:
+      return []
+
+    file_list = []
+    process_units = getattr(table, field_name)
+    # If the process_units field is not populated, it will be None. Use an
+    # empty list to skip the check.
+    for process_unit in process_units or []:
+      options = process_unit.options
+      if isinstance(options, (_metadata_fb.BertTokenizerOptionsT,
+                              _metadata_fb.RegexTokenizerOptionsT)):
+        file_list += self._get_associated_files_from_table(options, "vocabFile")
+      elif isinstance(options, _metadata_fb.SentencePieceTokenizerOptionsT):
+        file_list += self._get_associated_files_from_table(
+            options, "sentencePieceModel")
+        file_list += self._get_associated_files_from_table(options, "vocabFile")
+    return file_list
+
+  def _get_associated_files_from_table(self, table, field_name):
+    """Gets the associated files that are attached a table directly.
+
+    Args:
+      table: a Flatbuffers table object that contains fields of an array of
+        AssociatedFile, such as TensorMetadata and BertTokenizerOptions.
+      field_name: the name of the field in the table that represents an array of
+        ProcessUnit. If the table is TensorMetadata, field_name can be
+        "AssociatedFiles". If the table is BertTokenizerOptions, field_name can
+        be "VocabFile".
+
+    Returns:
+      A list of AssociatedFileT objects.
+    """
+
+    if table is None:
+      return []
+
+    # If the associated file field is not populated,
+    # `getattr(table, field_name)` will be None. Return an empty list.
+    return getattr(table, field_name) or []
+
+  def _get_recorded_associated_file_object_list(self, metadata):
+    """Gets a list of AssociatedFileT objects recorded in the metadata.
+
+    Associated files may be attached to a model, a subgraph, or an input/output
+    tensor.
+
+    Args:
+      metadata: the ModelMetadataT object.
+
+    Returns:
+      List of recorded AssociatedFileT objects.
+    """
+    recorded_files = []
+
+    # Add associated files attached to ModelMetadata.
+    recorded_files += self._get_associated_files_from_table(
+        metadata, "associatedFiles")
+
+    # Add associated files attached to each SubgraphMetadata.
+    for subgraph in metadata.subgraphMetadata or []:
+      recorded_files += self._get_associated_files_from_table(
+          subgraph, "associatedFiles")
+
+      # Add associated files attached to each input tensor.
+      for tensor_metadata in subgraph.inputTensorMetadata or []:
+        recorded_files += self._get_associated_files_from_table(
+            tensor_metadata, "associatedFiles")
+        recorded_files += self._get_associated_files_from_process_units(
+            tensor_metadata, "processUnits")
+
+      # Add associated files attached to each output tensor.
+      for tensor_metadata in subgraph.outputTensorMetadata or []:
+        recorded_files += self._get_associated_files_from_table(
+            tensor_metadata, "associatedFiles")
+        recorded_files += self._get_associated_files_from_process_units(
+            tensor_metadata, "processUnits")
+
+      # Add associated files attached to the input_process_units.
+      recorded_files += self._get_associated_files_from_process_units(
+          subgraph, "inputProcessUnits")
+
+      # Add associated files attached to the output_process_units.
+      recorded_files += self._get_associated_files_from_process_units(
+          subgraph, "outputProcessUnits")
+
+    return recorded_files
+
+  def _populate_associated_files(self):
+    """Concatenates associated files after TensorFlow Lite model file.
+
+    If the MetadataPopulator object is created using the method,
+    with_model_file(model_file), the model file will be updated.
+    """
+    # Opens up the model file in "appending" mode.
+    # If self._model_file already has pack files, zipfile will concatenate
+    # addition files after self._model_file. For example, suppose we have
+    # self._model_file = old_tflite_file | label1.txt | label2.txt
+    # Then after trigger populate() to add label3.txt, self._model_file becomes
+    # self._model_file = old_tflite_file | label1.txt | label2.txt | label3.txt
+    with tempfile.SpooledTemporaryFile() as temp:
+      # (1) Copy content from model file of to temp file.
+      with _open_file(self._model_file, "rb") as f:
+        shutil.copyfileobj(f, temp)
+
+      # (2) Append of to a temp file as a zip.
+      with _open_as_zipfile(temp, "a") as zf:
+        for file_name, file_buffer in self._associated_files.items():
+          zf.writestr(file_name, file_buffer)
+
+      # (3) Copy temp file to model file.
+      temp.seek(0)
+      with _open_file(self._model_file, "wb") as f:
+        shutil.copyfileobj(temp, f)
+
+  def _populate_metadata_buffer(self):
+    """Populates the metadata buffer (in bytearray) into the model file.
+
+    Inserts metadata_buf into the metadata field of schema.Model. If the
+    MetadataPopulator object is created using the method,
+    with_model_file(model_file), the model file will be updated.
+
+    Existing metadata buffer (if applied) will be overridden by the new metadata
+    buffer.
+    """
+
+    with _open_file(self._model_file, "rb") as f:
+      model_buf = f.read()
+
+    model = _schema_fb.ModelT.InitFromObj(
+        _schema_fb.Model.GetRootAsModel(model_buf, 0))
+    buffer_field = _schema_fb.BufferT()
+    buffer_field.data = self._metadata_buf
+
+    is_populated = False
+    if not model.metadata:
+      model.metadata = []
+    else:
+      # Check if metadata has already been populated.
+      for meta in model.metadata:
+        if meta.name.decode("utf-8") == self.METADATA_FIELD_NAME:
+          is_populated = True
+          model.buffers[meta.buffer] = buffer_field
+
+    if not is_populated:
+      if not model.buffers:
+        model.buffers = []
+      model.buffers.append(buffer_field)
+      # Creates a new metadata field.
+      metadata_field = _schema_fb.MetadataT()
+      metadata_field.name = self.METADATA_FIELD_NAME
+      metadata_field.buffer = len(model.buffers) - 1
+      model.metadata.append(metadata_field)
+
+    # Packs model back to a flatbuffer binaray file.
+    b = flatbuffers.Builder(0)
+    b.Finish(model.Pack(b), self.TFLITE_FILE_IDENTIFIER)
+    model_buf = b.Output()
+
+    # Saves the updated model buffer to model file.
+    # Gets files that have been packed to self._model_file.
+    packed_files = self.get_packed_associated_file_list()
+    if packed_files:
+      # Writes the updated model buffer and associated files into a new model
+      # file (in memory). Then overwrites the original model file.
+      with tempfile.SpooledTemporaryFile() as temp:
+        temp.write(model_buf)
+        self._copy_archived_files(self._model_file, packed_files, temp)
+        temp.seek(0)
+        with _open_file(self._model_file, "wb") as f:
+          shutil.copyfileobj(temp, f)
+    else:
+      with _open_file(self._model_file, "wb") as f:
+        f.write(model_buf)
+
+  def _use_basename_for_associated_files_in_metadata(self, metadata):
+    """Removes any associated file local directory (if exists)."""
+    for file in self._get_recorded_associated_file_object_list(metadata):
+      file.name = os.path.basename(file.name)
+
+  def _validate_metadata(self, metadata_buf):
+    """Validates the metadata to be populated."""
+    _assert_metadata_buffer_identifier(metadata_buf)
+
+    # Verify the number of SubgraphMetadata is exactly one.
+    # TFLite currently only support one subgraph.
+    model_meta = _metadata_fb.ModelMetadata.GetRootAsModelMetadata(
+        metadata_buf, 0)
+    if model_meta.SubgraphMetadataLength() != 1:
+      raise ValueError("The number of SubgraphMetadata should be exactly one, "
+                       "but got {0}.".format(
+                           model_meta.SubgraphMetadataLength()))
+
+    # Verify if the number of tensor metadata matches the number of tensors.
+    with _open_file(self._model_file, "rb") as f:
+      model_buf = f.read()
+    model = _schema_fb.Model.GetRootAsModel(model_buf, 0)
+
+    num_input_tensors = model.Subgraphs(0).InputsLength()
+    num_input_meta = model_meta.SubgraphMetadata(0).InputTensorMetadataLength()
+    if num_input_tensors != num_input_meta:
+      raise ValueError(
+          "The number of input tensors ({0}) should match the number of "
+          "input tensor metadata ({1})".format(num_input_tensors,
+                                               num_input_meta))
+    num_output_tensors = model.Subgraphs(0).OutputsLength()
+    num_output_meta = model_meta.SubgraphMetadata(
+        0).OutputTensorMetadataLength()
+    if num_output_tensors != num_output_meta:
+      raise ValueError(
+          "The number of output tensors ({0}) should match the number of "
+          "output tensor metadata ({1})".format(num_output_tensors,
+                                                num_output_meta))
+
+
+class _MetadataPopulatorWithBuffer(MetadataPopulator):
+  """Subclass of MetadataPopulator that populates metadata to a model buffer.
+
+  This class is used to populate metadata into a in-memory model buffer. As we
+  use Zip API to concatenate associated files after tflite model file, the
+  populating operation is developed based on a model file. For in-memory model
+  buffer, we create a tempfile to serve the populating operation. This class is
+  then used to generate this tempfile, and delete the file when the
+  MetadataPopulator object is deleted.
+  """
+
+  def __init__(self, model_buf):
+    """Constructor for _MetadataPopulatorWithBuffer.
+
+    Args:
+      model_buf: TensorFlow Lite model buffer in bytearray.
+
+    Raises:
+      ValueError: model_buf is empty.
+      ValueError: model_buf does not have the expected flatbuffer identifier.
+    """
+    if not model_buf:
+      raise ValueError("model_buf cannot be empty.")
+
+    with tempfile.NamedTemporaryFile() as temp:
+      model_file = temp.name
+
+    with _open_file(model_file, "wb") as f:
+      f.write(model_buf)
+
+    super().__init__(model_file)
+
+  def __del__(self):
+    """Destructor of _MetadataPopulatorWithBuffer.
+
+    Deletes the tempfile.
+    """
+    if os.path.exists(self._model_file):
+      os.remove(self._model_file)
+
+
+class MetadataDisplayer(object):
+  """Displays metadata and associated file info in human-readable format."""
+
+  def __init__(self, model_buffer, metadata_buffer, associated_file_list):
+    """Constructor for MetadataDisplayer.
+
+    Args:
+      model_buffer: valid buffer of the model file.
+      metadata_buffer: valid buffer of the metadata file.
+      associated_file_list: list of associate files in the model file.
+    """
+    _assert_model_buffer_identifier(model_buffer)
+    _assert_metadata_buffer_identifier(metadata_buffer)
+    self._model_buffer = model_buffer
+    self._metadata_buffer = metadata_buffer
+    self._associated_file_list = associated_file_list
+
+  @classmethod
+  def with_model_file(cls, model_file):
+    """Creates a MetadataDisplayer object for the model file.
+
+    Args:
+      model_file: valid path to a TensorFlow Lite model file.
+
+    Returns:
+      MetadataDisplayer object.
+
+    Raises:
+      IOError: File not found.
+      ValueError: The model does not have metadata.
+    """
+    _assert_file_exist(model_file)
+    with _open_file(model_file, "rb") as f:
+      return cls.with_model_buffer(f.read())
+
+  @classmethod
+  def with_model_buffer(cls, model_buffer):
+    """Creates a MetadataDisplayer object for a file buffer.
+
+    Args:
+      model_buffer: TensorFlow Lite model buffer in bytearray.
+
+    Returns:
+      MetadataDisplayer object.
+    """
+    if not model_buffer:
+      raise ValueError("model_buffer cannot be empty.")
+    metadata_buffer = get_metadata_buffer(model_buffer)
+    if not metadata_buffer:
+      raise ValueError("The model does not have metadata.")
+    associated_file_list = cls._parse_packed_associated_file_list(model_buffer)
+    return cls(model_buffer, metadata_buffer, associated_file_list)
+
+  def get_associated_file_buffer(self, filename):
+    """Get the specified associated file content in bytearray.
+
+    Args:
+      filename: name of the file to be extracted.
+
+    Returns:
+      The file content in bytearray.
+
+    Raises:
+      ValueError: if the file does not exist in the model.
+    """
+    if filename not in self._associated_file_list:
+      raise ValueError(
+          "The file, {}, does not exist in the model.".format(filename))
+
+    with _open_as_zipfile(io.BytesIO(self._model_buffer)) as zf:
+      return zf.read(filename)
+
+  def get_metadata_buffer(self):
+    """Get the metadata buffer in bytearray out from the model."""
+    return copy.deepcopy(self._metadata_buffer)
+
+  def get_metadata_json(self):
+    """Converts the metadata into a json string."""
+    return convert_to_json(self._metadata_buffer)
+
+  def get_packed_associated_file_list(self):
+    """Returns a list of associated files that are packed in the model.
+
+    Returns:
+      A name list of associated files.
+    """
+    return copy.deepcopy(self._associated_file_list)
+
+  @classmethod
+  def _parse_packed_associated_file_list(cls, model_buf):
+    """Gets a list of associated files packed to the model file.
+
+    Args:
+      model_buf: valid file buffer.
+
+    Returns:
+      List of packed associated files.
+    """
+
+    try:
+      with _open_as_zipfile(io.BytesIO(model_buf)) as zf:
+        return zf.namelist()
+    except zipfile.BadZipFile:
+      return []
+
+
+def _get_custom_metadata(metadata_buffer: bytes, name: str):
+  """Gets the custom metadata in metadata_buffer based on the name.
+
+  Args:
+    metadata_buffer: valid metadata buffer in bytes.
+    name: custom metadata name.
+
+  Returns:
+    Index of custom metadata, custom metadata flatbuffer. Returns (None, None)
+    if the custom metadata is not found.
+  """
+  model_metadata = _metadata_fb.ModelMetadata.GetRootAs(metadata_buffer)
+  subgraph = model_metadata.SubgraphMetadata(0)
+  if subgraph is None or subgraph.CustomMetadataIsNone():
+    return None, None
+
+  for i in range(subgraph.CustomMetadataLength()):
+    custom_metadata = subgraph.CustomMetadata(i)
+    if custom_metadata.Name().decode("utf-8") == name:
+      return i, custom_metadata.DataAsNumpy().tobytes()
+  return None, None
+
+
+# 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, custom_metadata_schema: Optional[Dict[str, str]] = None
+) -> str:
+  """Converts the metadata into a json string.
+
+  Args:
+    metadata_buffer: valid metadata buffer in bytes.
+    custom_metadata_schema: A dict of custom metadata schema, in which key is
+      custom metadata name [1], value is the filepath that defines custom
+      metadata schema. For instance, custom_metadata_schema =
+      {"SEGMENTER_METADATA": "metadata/vision_tasks_metadata_schema.fbs"}. [1]:
+        https://github.com/google/mediapipe/blob/46b5c4012d2ef76c9d92bb0d88a6b107aee83814/mediapipe/tasks/metadata/metadata_schema.fbs#L612
+
+  Returns:
+    Metadata in JSON format.
+
+  Raises:
+    ValueError: error occurred when parsing the metadata schema file.
+  """
+  opt = _pywrap_flatbuffers.IDLOptions()
+  opt.strict_json = True
+  parser = _pywrap_flatbuffers.Parser(opt)
+  with _open_file(_FLATC_TFLITE_METADATA_SCHEMA_FILE) as f:
+    metadata_schema_content = f.read()
+  if not parser.parse(metadata_schema_content):
+    raise ValueError("Cannot parse metadata schema. Reason: " + parser.error)
+  # Json content which may contain binary custom metadata.
+  raw_json_content = _pywrap_flatbuffers.generate_text(parser, metadata_buffer)
+  if not custom_metadata_schema:
+    return raw_json_content
+
+  json_data = json.loads(raw_json_content)
+  # Gets the custom metadata by name and parse the binary custom metadata into
+  # human readable json content.
+  for name, schema_file in custom_metadata_schema.items():
+    idx, custom_metadata = _get_custom_metadata(metadata_buffer, name)
+    if not custom_metadata:
+      logging.info(
+          "No custom metadata with name %s in metadata flatbuffer.", name
+      )
+      continue
+    _assert_file_exist(schema_file)
+    with _open_file(schema_file, "rb") as f:
+      custom_metadata_schema_content = f.read()
+    if not parser.parse(custom_metadata_schema_content):
+      raise ValueError(
+          "Cannot parse custom metadata schema. Reason: " + parser.error
+      )
+    custom_metadata_json = _pywrap_flatbuffers.generate_text(
+        parser, custom_metadata
+    )
+    json_meta = json_data["subgraph_metadata"][0]["custom_metadata"][idx]
+    json_meta["name"] = name
+    json_meta["data"] = json.loads(custom_metadata_json)
+  return json.dumps(json_data, indent=2)
+
+
+def _assert_file_exist(filename):
+  """Checks if a file exists."""
+  if not _exists_file(filename):
+    raise IOError("File, '{0}', does not exist.".format(filename))
+
+
+def _assert_model_file_identifier(model_file):
+  """Checks if a model file has the expected TFLite schema identifier."""
+  _assert_file_exist(model_file)
+  with _open_file(model_file, "rb") as f:
+    _assert_model_buffer_identifier(f.read())
+
+
+def _assert_model_buffer_identifier(model_buf):
+  if not _schema_fb.Model.ModelBufferHasIdentifier(model_buf, 0):
+    raise ValueError(
+        "The model provided does not have the expected identifier, and "
+        "may not be a valid TFLite model.")
+
+
+def _assert_metadata_buffer_identifier(metadata_buf):
+  """Checks if a metadata buffer has the expected Metadata schema identifier."""
+  if not _metadata_fb.ModelMetadata.ModelMetadataBufferHasIdentifier(
+      metadata_buf, 0):
+    raise ValueError(
+        "The metadata buffer does not have the expected identifier, and may not"
+        " be a valid TFLite Metadata.")
+
+
+def get_metadata_buffer(model_buf):
+  """Returns the metadata in the model file as a buffer.
+
+  Args:
+    model_buf: valid buffer of the model file.
+
+  Returns:
+    Metadata buffer. Returns `None` if the model does not have metadata.
+  """
+  tflite_model = _schema_fb.Model.GetRootAsModel(model_buf, 0)
+
+  # Gets metadata from the model file.
+  for i in range(tflite_model.MetadataLength()):
+    meta = tflite_model.Metadata(i)
+    if meta.Name().decode("utf-8") == MetadataPopulator.METADATA_FIELD_NAME:
+      buffer_index = meta.Buffer()
+      metadata = tflite_model.Buffers(buffer_index)
+      if metadata.DataLength() == 0:
+        continue
+      return metadata.DataAsNumpy().tobytes()
+
+  return None