Internal change

PiperOrigin-RevId: 481515490

mediapipe/tasks/python/metadata/metadata_writers/BUILD

@@ -19,3 +19,21 @@ py_library(
         "//mediapipe/tasks/metadata:schema_py",
     ],
 )
+
+py_library(
+    name = "metadata_writer",
+    srcs = ["metadata_writer.py"],
+    deps = [
+        ":metadata_info",
+        ":writer_utils",
+        "//mediapipe/tasks/metadata:metadata_schema_py",
+        "//mediapipe/tasks/python/metadata",
+        "@flatbuffers//:runtime_py",
+    ],
+)
+
+py_library(
+    name = "writer_utils",
+    srcs = ["writer_utils.py"],
+    deps = ["//mediapipe/tasks/metadata:schema_py"],
+)

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

@@ -0,0 +1,468 @@
+# Copyright 2022 The MediaPipe Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Generic metadata writer."""
+
+import collections
+import dataclasses
+import os
+import tempfile
+from typing import List, Optional, Tuple
+
+import flatbuffers
+from mediapipe.tasks.metadata import metadata_schema_py_generated as _metadata_fb
+from mediapipe.tasks.python.metadata import metadata as _metadata
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_info
+from mediapipe.tasks.python.metadata.metadata_writers import writer_utils
+
+_INPUT_IMAGE_NAME = 'image'
+_INPUT_IMAGE_DESCRIPTION = 'Input image to be processed.'
+
+
+@dataclasses.dataclass
+class CalibrationParameter:
+  """Parameters for score calibration [1].
+
+  Score calibration is performed on an output tensor through sigmoid functions.
+  One of 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.
+
+  For each index in the output tensor, this applies:
+    * `f(x) = scale / (1 + e^-(slope * g(x) + offset))` if `x > min_score` or if
+      no `min_score` has been specified.
+    * `f(x) = default_score` otherwise or if no scale, slope and offset have
+      been specified.
+
+  [1]:
+    https://github.com/google/mediapipe/blob/f8af41b1eb49ff4bdad756ff19d1d36f486be614/mediapipe/tasks/metadata/metadata_schema.fbs#L434
+  """
+  scale: float
+  slope: float
+  offset: float
+  min_score: Optional[float] = None
+
+
+@dataclasses.dataclass
+class LabelItem:
+  """Label item for labels per locale.
+
+  Attributes:
+    filename: The file name to save the labels.
+    names: A list of label names.
+    locale: The specified locale for labels.
+  """
+  filename: str
+  names: List[str]
+  locale: Optional[str] = None
+
+
+class Labels(object):
+  """Simple container holding classification labels of a particular tensor.
+
+  Example usage:
+    # The first added label list can be used as category names as needed.
+    labels = Labels()
+      .add(['/m/011l78', '/m/031d23'])
+      .add(['cat', 'dog], 'en')
+      .add(['chat', 'chien], 'fr')
+  """
+
+  def __init__(self) -> None:
+    self._labels = []  # [LabelItem]
+
+  @property
+  def labels(self) -> List[LabelItem]:
+    return self._labels
+
+  def add(self,
+          labels: List[str],
+          locale: Optional[str] = None,
+          exported_filename: Optional[str] = None) -> 'Labels':
+    """Adds labels in the container.
+
+    Args:
+      labels: A list of label names, e.g. ['apple', 'pear', 'banana'].
+      locale: The specified locale for labels.
+      exported_filename: The file name to export the labels. If not set,
+        filename defaults to 'labels.txt'.
+
+    Returns:
+      The Labels instance, can be used for chained operation.
+    """
+    if not labels:
+      raise ValueError('The list of labels is empty.')
+
+    # Prepare the new item to be inserted
+    if not exported_filename:
+      exported_filename = 'labels'
+      if locale:
+        exported_filename += f'_{locale}'
+      exported_filename += '.txt'
+    item = LabelItem(filename=exported_filename, names=labels, locale=locale)
+
+    # Insert the new element at the end of the list
+    self._labels.append(item)
+    return self
+
+  def add_from_file(self,
+                    label_filepath: str,
+                    locale: Optional[str] = None,
+                    exported_filename: Optional[str] = None) -> 'Labels':
+    """Adds a label file in the container.
+
+    Args:
+      label_filepath: File path to read labels. Each line is a label name in the
+        file.
+      locale: The specified locale for labels.
+      exported_filename: The file name to export the labels. If not set,
+        filename defaults to 'labels.txt'.
+
+    Returns:
+      The Labels instance, can be used for chained operation.
+    """
+
+    with open(label_filepath, 'r') as f:
+      labels = f.read().split('\n')
+      return self.add(labels, locale, exported_filename)
+
+
+class ScoreCalibration:
+  """Simple container holding score calibration related parameters."""
+
+  # A shortcut to avoid client side code importing _metadata_fb
+  transformation_types = _metadata_fb.ScoreTransformationType
+
+  def __init__(self,
+               transformation_type: _metadata_fb.ScoreTransformationType,
+               parameters: List[CalibrationParameter],
+               default_score: int = 0):
+    self.transformation_type = transformation_type
+    self.parameters = parameters
+    self.default_score = default_score
+
+
+def _fill_default_tensor_names(
+    tensor_metadata: List[_metadata_fb.TensorMetadataT],
+    tensor_names_from_model: List[str]):
+  """Fills the default tensor names."""
+  # If tensor name in metadata is empty, default to the tensor name saved in
+  # the model.
+  for metadata, name in zip(tensor_metadata, tensor_names_from_model):
+    metadata.name = metadata.name or name
+
+
+def _pair_tensor_metadata(
+    tensor_md: List[metadata_info.TensorMd],
+    tensor_names_from_model: List[str]) -> List[metadata_info.TensorMd]:
+  """Pairs tensor_md according to the tensor names from the model."""
+  tensor_names_from_arg = [
+      md.tensor_name for md in tensor_md or [] if md.tensor_name is not None
+  ]
+  if not tensor_names_from_arg:
+    return tensor_md
+
+  if collections.Counter(tensor_names_from_arg) != collections.Counter(
+      tensor_names_from_model):
+    raise ValueError(
+        'The tensor names from arguments ({}) do not match the tensor names'
+        ' read from the model ({}).'.format(tensor_names_from_arg,
+                                            tensor_names_from_model))
+  pairs_tensor_md = []
+  name_md_dict = dict(zip(tensor_names_from_arg, tensor_md))
+  for name in tensor_names_from_model:
+    pairs_tensor_md.append(name_md_dict[name])
+  return pairs_tensor_md
+
+
+def _create_metadata_buffer(
+    model_buffer: bytearray,
+    general_md: Optional[metadata_info.GeneralMd] = None,
+    input_md: Optional[List[metadata_info.TensorMd]] = None,
+    output_md: Optional[List[metadata_info.TensorMd]] = None) -> bytearray:
+  """Creates a buffer of the metadata.
+
+  Args:
+    model_buffer: valid buffer of the model file.
+    general_md: general information about the model.
+    input_md: metadata information of the input tensors.
+    output_md: metadata information of the output tensors.
+
+  Returns:
+    A buffer of the metadata.
+
+  Raises:
+    ValueError: if the tensor names from `input_md` and `output_md` do not
+    match the tensor names read from the model.
+  """
+  # Create input metadata from `input_md`.
+  if input_md:
+    input_md = _pair_tensor_metadata(
+        input_md, writer_utils.get_input_tensor_names(model_buffer))
+    input_metadata = [m.create_metadata() for m in input_md]
+  else:
+    num_input_tensors = writer_utils.get_subgraph(model_buffer).InputsLength()
+    input_metadata = [_metadata_fb.TensorMetadataT()] * num_input_tensors
+
+  _fill_default_tensor_names(input_metadata,
+                             writer_utils.get_input_tensor_names(model_buffer))
+
+  # Create output metadata from `output_md`.
+  if output_md:
+    output_md = _pair_tensor_metadata(
+        output_md, writer_utils.get_output_tensor_names(model_buffer))
+    output_metadata = [m.create_metadata() for m in output_md]
+  else:
+    num_output_tensors = writer_utils.get_subgraph(model_buffer).OutputsLength()
+    output_metadata = [_metadata_fb.TensorMetadataT()] * num_output_tensors
+  _fill_default_tensor_names(output_metadata,
+                             writer_utils.get_output_tensor_names(model_buffer))
+
+  # Create the subgraph metadata.
+  subgraph_metadata = _metadata_fb.SubGraphMetadataT()
+  subgraph_metadata.inputTensorMetadata = input_metadata
+  subgraph_metadata.outputTensorMetadata = output_metadata
+
+  # Create the whole model metadata.
+  if general_md is None:
+    general_md = metadata_info.GeneralMd()
+  model_metadata = general_md.create_metadata()
+  model_metadata.subgraphMetadata = [subgraph_metadata]
+
+  # Get the metadata flatbuffer.
+  b = flatbuffers.Builder(0)
+  b.Finish(
+      model_metadata.Pack(b),
+      _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
+  return b.Output()
+
+
+class MetadataWriter(object):
+  """Generic Metadata writer.
+
+  Example usage:
+
+  For an example model which requires two inputs: image and general feature
+  inputs, and generates one output: classification.
+
+  with open(model_path, 'rb') as f:
+    writer = MetadataWriter.create(f.read())
+    model_content, metadata_json_content = writer
+        .add_genernal_info('model_name', 'model description')
+        .add_image_input()
+        .add_feature_input()
+        .add_classification_output(Labels().add(['A', 'B']))
+        .populate()
+  """
+
+  @classmethod
+  def create(cls, model_buffer: bytearray) -> 'MetadataWriter':
+    return cls(model_buffer)
+
+  def __init__(self, model_buffer: bytearray) -> None:
+    self._model_buffer = model_buffer
+    self._general_md = None
+    self._input_mds = []
+    self._output_mds = []
+    self._associated_files = []
+    self._temp_folder = tempfile.TemporaryDirectory()
+
+  def __del__(self):
+    if os.path.exists(self._temp_folder.name):
+      self._temp_folder.cleanup()
+
+  def add_genernal_info(
+      self,
+      model_name: str,
+      model_description: Optional[str] = None) -> 'MetadataWriter':
+    """Adds a genernal info metadata for the general metadata informantion."""
+    # Will overwrite the previous `self._general_md` if exists.
+    self._general_md = metadata_info.GeneralMd(
+        name=model_name, description=model_description)
+    return self
+
+  color_space_types = _metadata_fb.ColorSpaceType
+
+  def add_feature_input(self,
+                        name: Optional[str] = None,
+                        description: Optional[str] = None) -> 'MetadataWriter':
+    """Adds an input tensor metadata for the general basic feature input."""
+    input_md = metadata_info.TensorMd(name=name, description=description)
+    self._input_mds.append(input_md)
+    return self
+
+  def add_image_input(
+      self,
+      norm_mean: List[float],
+      norm_std: List[float],
+      color_space_type: Optional[int] = _metadata_fb.ColorSpaceType.RGB,
+      name: str = _INPUT_IMAGE_NAME,
+      description: str = _INPUT_IMAGE_DESCRIPTION) -> 'MetadataWriter':
+    """Adds an input image metadata for the image input.
+
+    Args:
+      norm_mean: The mean value used to normalize each input channel. If there
+        is only one element in the list, its value will be broadcasted to all
+        channels. Also note that norm_mean and norm_std should have the same
+        number of elements. [1]
+      norm_std: The std value used to normalize each input channel. If there is
+        only one element in the list, its value will be broadcasted to all
+        channels. [1]
+      color_space_type: The color space type of the input image. [2]
+      name: Name of the input tensor.
+      description: Description of the input tensor.
+
+    Returns:
+      The MetaWriter instance, can be used for chained operation.
+
+    [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
+    """
+    input_md = metadata_info.InputImageTensorMd(
+        name=name,
+        description=description,
+        norm_mean=norm_mean,
+        norm_std=norm_std,
+        color_space_type=color_space_type,
+        tensor_type=self._input_tensor_type(len(self._input_mds)))
+
+    self._input_mds.append(input_md)
+    return self
+
+  _OUTPUT_CLASSIFICATION_NAME = 'score'
+  _OUTPUT_CLASSIFICATION_DESCRIPTION = 'Score of the labels respectively'
+
+  def add_classification_output(
+      self,
+      labels: Optional[Labels] = None,
+      score_calibration: Optional[ScoreCalibration] = None,
+      name: str = _OUTPUT_CLASSIFICATION_NAME,
+      description: str = _OUTPUT_CLASSIFICATION_DESCRIPTION
+  ) -> 'MetadataWriter':
+    """Add a classification head metadata for classification output tensor.
+
+    Example usage:
+      writer.add_classification_output(
+        Labels()
+          .add(['/m/011l78', '/m/031d23'])
+          .add(['cat', 'dog], 'en')
+          .add(['chat', 'chien], 'fr')
+          )
+
+    Args:
+      labels: an instance of Labels helper class.
+      score_calibration: an instance of ScoreCalibration helper class.
+      name: Metadata name of the tensor. Note that this is different from tensor
+        name in the flatbuffer.
+      description: human readable description of what the output is.
+
+    Returns:
+      The current Writer instance to allow chained operation.
+    """
+    calibration_md = None
+    if score_calibration:
+      calibration_md = metadata_info.ScoreCalibrationMd(
+          score_transformation_type=score_calibration.transformation_type,
+          default_score=score_calibration.default_score,
+          file_path=self._export_calibration_file('score_calibration.txt',
+                                                  score_calibration.parameters))
+
+    label_files = None
+    if labels:
+      label_files = []
+      for item in labels.labels:
+        label_files.append(
+            metadata_info.LabelFileMd(
+                self._export_labels(item.filename, item.names),
+                locale=item.locale))
+
+    output_md = metadata_info.ClassificationTensorMd(
+        name=name,
+        description=description,
+        label_files=label_files,
+        tensor_type=self._output_tensor_type(len(self._output_mds)),
+        score_calibration_md=calibration_md,
+    )
+    self._output_mds.append(output_md)
+    return self
+
+  def add_feature_output(self,
+                         name: Optional[str] = None,
+                         description: Optional[str] = None) -> 'MetadataWriter':
+    """Adds an output tensor metadata for the general basic feature output."""
+    output_md = metadata_info.TensorMd(name=name, description=description)
+    self._output_mds.append(output_md)
+    return self
+
+  def populate(self) -> Tuple[bytearray, str]:
+    """Populates metadata into the TFLite file.
+
+    Note that only the output tflite is used for deployment. The output JSON
+    content is used to interpret the metadata content.
+
+    Returns:
+      A tuple of (model_with_metadata_in_bytes, metdata_json_content)
+    """
+    # Populates metadata and associated files into TFLite model buffer.
+    populator = _metadata.MetadataPopulator.with_model_buffer(
+        self._model_buffer)
+    metadata_buffer = _create_metadata_buffer(
+        model_buffer=self._model_buffer,
+        general_md=self._general_md,
+        input_md=self._input_mds,
+        output_md=self._output_mds)
+    populator.load_metadata_buffer(metadata_buffer)
+    if self._associated_files:
+      populator.load_associated_files(self._associated_files)
+    populator.populate()
+    tflite_content = populator.get_model_buffer()
+
+    displayer = _metadata.MetadataDisplayer.with_model_buffer(tflite_content)
+    metadata_json_content = displayer.get_metadata_json()
+
+    return tflite_content, metadata_json_content
+
+  def _input_tensor_type(self, idx):
+    return writer_utils.get_input_tensor_types(self._model_buffer)[idx]
+
+  def _output_tensor_type(self, idx):
+    return writer_utils.get_output_tensor_types(self._model_buffer)[idx]
+
+  def _export_labels(self, filename: str, index_to_label: List[str]) -> str:
+    filepath = os.path.join(self._temp_folder.name, filename)
+    with open(filepath, 'w') as f:
+      f.write('\n'.join(index_to_label))
+    self._associated_files.append(filepath)
+    return filepath
+
+  def _export_calibration_file(self, filename: str,
+                               calibrations: List[CalibrationParameter]) -> str:
+    """Stores calibration parameters in a csv file."""
+    filepath = os.path.join(self._temp_folder.name, filename)
+    with open(filepath, 'w') as f:
+      for idx, item in enumerate(calibrations):
+        if idx != 0:
+          f.write('\n')
+        if item:
+          if item.scale is None or item.slope is None or item.offset is None:
+            raise ValueError('scale, slope and offset values can not be set to '
+                             'None.')
+          elif item.min_score is not None:
+            f.write(f'{item.scale},{item.slope},{item.offset},{item.min_score}')
+          else:
+            f.write(f'{item.scale},{item.slope},{item.offset}')
+
+          self._associated_files.append(filepath)
+    return filepath

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

@@ -0,0 +1,85 @@
+# Copyright 2022 The MediaPipe Authors. All Rights Reserved.
+#
+# 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_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/metadata/metadata_writers/BUILD

@@ -24,3 +24,13 @@ py_test(
         "@flatbuffers//:runtime_py",
     ],
 )
+
+py_test(
+    name = "metadata_writer_test",
+    srcs = ["metadata_writer_test.py"],
+    data = ["//mediapipe/tasks/testdata/metadata:model_files"],
+    deps = [
+        "//mediapipe/tasks/python/metadata/metadata_writers:metadata_writer",
+        "//mediapipe/tasks/python/test:test_utils",
+    ],
+)

mediapipe/tasks/python/test/metadata/metadata_writers/metadata_writer_test.py

@@ -0,0 +1,358 @@
+# Copyright 2022 The MediaPipe Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for metadata writer classes."""
+from absl.testing import absltest
+
+from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
+from mediapipe.tasks.python.test import test_utils
+
+_IMAGE_CLASSIFIER_MODEL = test_utils.get_test_data_path(
+    'mobilenet_v1_0.25_224_1_default_1.tflite')
+
+
+class LabelsTest(absltest.TestCase):
+
+  def test_category_name(self):
+    labels = metadata_writer.Labels()
+    self.assertEqual(
+        labels.add(['a', 'b'])._labels, [
+            metadata_writer.LabelItem(
+                filename='labels.txt', names=['a', 'b'], locale=None)
+        ])
+
+  def test_locale(self):
+    labels = metadata_writer.Labels()
+
+    # Add from file.
+    en_filepath = self.create_tempfile().full_path
+    with open(en_filepath, 'w') as f:
+      f.write('a\nb')
+    labels.add_from_file(en_filepath, 'en')
+
+    # Customized file name.
+    labels.add(['A', 'B'], 'fr', exported_filename='my_file.txt')
+    self.assertEqual(labels._labels, [
+        metadata_writer.LabelItem('labels_en.txt', ['a', 'b'], 'en'),
+        metadata_writer.LabelItem('my_file.txt', ['A', 'B'], 'fr'),
+    ])
+
+
+class MetadataWriterForTaskTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    with open(_IMAGE_CLASSIFIER_MODEL, 'rb') as f:
+      self.image_classifier_model_buffer = f.read()
+
+  def test_initialize_and_populate(self):
+    writer = metadata_writer.MetadataWriter.create(
+        self.image_classifier_model_buffer)
+    writer.add_genernal_info(
+        model_name='my_image_model', model_description='my_description')
+    tflite_model, metadata_json = writer.populate()
+    self.assertLen(tflite_model, 1882986)
+    self.assertJsonEqual(
+        metadata_json, """{
+          "name": "my_image_model",
+          "description": "my_description",
+          "subgraph_metadata": [
+            {
+              "input_tensor_metadata": [
+                {
+                  "name": "input"
+                }
+              ],
+              "output_tensor_metadata": [
+                {
+                  "name": "MobilenetV1/Predictions/Reshape_1"
+                }
+              ]
+            }
+          ],
+          "min_parser_version": "1.0.0"
+        }
+        """)
+
+  def test_add_feature_input_output(self):
+    writer = metadata_writer.MetadataWriter.create(
+        self.image_classifier_model_buffer)
+    writer.add_genernal_info(
+        model_name='my_model', model_description='my_description')
+    writer.add_feature_input(
+        name='input_tesnor', description='a feature input tensor')
+    writer.add_feature_output(
+        name='output_tesnor', description='a feature output tensor')
+
+    _, metadata_json = writer.populate()
+    self.assertJsonEqual(
+        metadata_json, """{
+          "name": "my_model",
+          "description": "my_description",
+          "subgraph_metadata": [
+            {
+              "input_tensor_metadata": [
+                {
+                  "name": "input_tesnor",
+                  "description": "a feature input tensor",
+                  "content": {
+                    "content_properties_type": "FeatureProperties",
+                    "content_properties": {
+                    }
+                  },
+                  "stats": {
+                  }
+                }
+              ],
+              "output_tensor_metadata": [
+                {
+                  "name": "output_tesnor",
+                  "description": "a feature output tensor",
+                  "content": {
+                    "content_properties_type": "FeatureProperties",
+                    "content_properties": {
+                    }
+                  },
+                  "stats": {
+                  }
+                }
+              ]
+            }
+          ],
+          "min_parser_version": "1.0.0"
+        }
+        """)
+
+  def test_image_classifier(self):
+    writer = metadata_writer.MetadataWriter.create(
+        self.image_classifier_model_buffer)
+    writer.add_genernal_info(
+        model_name='image_classifier',
+        model_description='Imagenet classification model')
+    writer.add_image_input(
+        norm_mean=[127.5, 127.5, 127.5],
+        norm_std=[127.5, 127.5, 127.5],
+        color_space_type=metadata_writer.MetadataWriter.color_space_types.RGB)
+    writer.add_classification_output(metadata_writer.Labels().add(
+        ['a', 'b', 'c']))
+    _, metadata_json = writer.populate()
+    self.assertJsonEqual(
+        metadata_json, """{
+          "name": "image_classifier",
+          "description": "Imagenet classification model",
+          "subgraph_metadata": [
+            {
+              "input_tensor_metadata": [
+                {
+                  "name": "image",
+                  "description": "Input image to be processed.",
+                  "content": {
+                    "content_properties_type": "ImageProperties",
+                    "content_properties": {
+                      "color_space": "RGB"
+                    }
+                  },
+                  "process_units": [
+                    {
+                      "options_type": "NormalizationOptions",
+                      "options": {
+                        "mean": [
+                          127.5,
+                          127.5,
+                          127.5
+                        ],
+                        "std": [
+                          127.5,
+                          127.5,
+                          127.5
+                        ]
+                      }
+                    }
+                  ],
+                  "stats": {
+                    "max": [
+                      1.0,
+                      1.0,
+                      1.0
+                    ],
+                    "min": [
+                      -1.0,
+                      -1.0,
+                      -1.0
+                    ]
+                  }
+                }
+              ],
+              "output_tensor_metadata": [
+                {
+                  "name": "score",
+                  "description": "Score of the labels respectively",
+                  "content": {
+                    "content_properties_type": "FeatureProperties",
+                    "content_properties": {
+                    }
+                  },
+                  "stats": {
+                    "max": [
+                      1.0
+                    ],
+                    "min": [
+                      0.0
+                    ]
+                  },
+                  "associated_files": [
+                    {
+                      "name": "labels.txt",
+                      "description": "Labels for categories that the model can recognize.",
+                      "type": "TENSOR_AXIS_LABELS"
+                    }
+                  ]
+                }
+              ]
+            }
+          ],
+          "min_parser_version": "1.0.0"
+        }
+        """)
+
+  def test_image_classifier_with_locale_and_score_calibration(self):
+    writer = metadata_writer.MetadataWriter(self.image_classifier_model_buffer)
+    writer.add_genernal_info(
+        model_name='image_classifier',
+        model_description='Classify the input image.')
+    writer.add_image_input(
+        norm_mean=[127.5, 127.5, 127.5],
+        norm_std=[127.5, 127.5, 127.5],
+        color_space_type=metadata_writer.MetadataWriter.color_space_types.RGB)
+    writer.add_classification_output(
+        metadata_writer.Labels().add(['/id1',
+                                      '/id2']).add(['tulip', 'lily'],
+                                                   'en').add(['tulipe', 'lis'],
+                                                             'fr'),
+        score_calibration=metadata_writer.ScoreCalibration(
+            metadata_writer.ScoreCalibration.transformation_types
+            .INVERSE_LOGISTIC, [
+                metadata_writer.CalibrationParameter(1., 2., 3., None),
+                metadata_writer.CalibrationParameter(1., 2., 3., 4.),
+            ],
+            default_score=0.5))
+    _, metadata_json = writer.populate()
+    self.assertJsonEqual(
+        metadata_json, """{
+          "name": "image_classifier",
+          "description": "Classify the input image.",
+          "subgraph_metadata": [
+            {
+              "input_tensor_metadata": [
+                {
+                  "name": "image",
+                  "description": "Input image to be processed.",
+                  "content": {
+                    "content_properties_type": "ImageProperties",
+                    "content_properties": {
+                      "color_space": "RGB"
+                    }
+                  },
+                  "process_units": [
+                    {
+                      "options_type": "NormalizationOptions",
+                      "options": {
+                        "mean": [
+                          127.5,
+                          127.5,
+                          127.5
+                        ],
+                        "std": [
+                          127.5,
+                          127.5,
+                          127.5
+                        ]
+                      }
+                    }
+                  ],
+                  "stats": {
+                    "max": [
+                      1.0,
+                      1.0,
+                      1.0
+                    ],
+                    "min": [
+                      -1.0,
+                      -1.0,
+                      -1.0
+                    ]
+                  }
+                }
+              ],
+              "output_tensor_metadata": [
+                {
+                  "name": "score",
+                  "description": "Score of the labels respectively",
+                  "content": {
+                    "content_properties_type": "FeatureProperties",
+                    "content_properties": {
+                    }
+                  },
+                  "process_units": [
+                    {
+                      "options_type": "ScoreCalibrationOptions",
+                      "options": {
+                        "score_transformation": "INVERSE_LOGISTIC",
+                        "default_score": 0.5
+                      }
+                    }
+                  ],
+                  "stats": {
+                    "max": [
+                      1.0
+                    ],
+                    "min": [
+                      0.0
+                    ]
+                  },
+                  "associated_files": [
+                    {
+                      "name": "labels.txt",
+                      "description": "Labels for categories that the model can recognize.",
+                      "type": "TENSOR_AXIS_LABELS"
+                    },
+                    {
+                      "name": "labels_en.txt",
+                      "description": "Labels for categories that the model can recognize.",
+                      "type": "TENSOR_AXIS_LABELS",
+                      "locale": "en"
+                    },
+                    {
+                      "name": "labels_fr.txt",
+                      "description": "Labels for categories that the model can recognize.",
+                      "type": "TENSOR_AXIS_LABELS",
+                      "locale": "fr"
+                    },
+                    {
+                      "name": "score_calibration.txt",
+                      "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.",
+                      "type": "TENSOR_AXIS_SCORE_CALIBRATION"
+                    }
+                  ]
+                }
+              ]
+            }
+          ],
+          "min_parser_version": "1.0.0"
+        }
+        """)
+
+
+if __name__ == '__main__':
+  absltest.main()