Internal change
PiperOrigin-RevId: 481515490
This commit is contained in:
Yuqi Li
Copybara-Service
parent
2def7c8203
commit
660f1812c2
|
|
@ -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"],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue
Block a user