def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[List[Type[metadata_info.TensorMd]]] = None,
output_md: Optional[List[Type[metadata_info.TensorMd]]] = None,
associated_files: Optional[List[str]] = None):
"""Creates MetadataWriter based on the metadata information.
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.
associated_files: path to the associated files to be populated.
Returns:
A MetadataWriter Object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd()
if input_md is None:
input_md = []
if output_md is None:
output_md = []
model_metadata = general_md.create_metadata()
input_metadata = [m.create_metadata() for m in input_md]
output_metadata = [m.create_metadata() for m in output_md]
return cls.create_from_metadata(model_buffer, model_metadata,
input_metadata, output_metadata,
associated_files)
def __init__(self, model_buffer: bytearray, model_name: str,
model_description: str):
self._model_buffer = model_buffer
self._general_md = metadata_info.GeneralMd(
name=model_name, description=model_description)
self._input_mds = []
self._output_mds = []
self._associate_files = []
def test_populate_create_from_metadata_info_should_succeed(self):
model_buffer = test_utils.load_file(_MODEL)
general_md = metadata_info.GeneralMd(name=_MODEL_NAME)
input_md = metadata_info.TensorMd(name=_INPUT_NAME)
output_md = metadata_info.TensorMd(name=_OUTPUT_NAME)
writer = metadata_writer.MetadataWriter.create_from_metadata_info(
model_buffer, general_md, [input_md], [output_md], [_LABEL_FILE])
model_with_metadata = writer.populate()
self._assert_correct_metadata(model_with_metadata,
_EXPECTED_META_INFO_JSON)
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputTextTensorMd] = None,
output_md: Optional[metadata_info.ClassificationTensorMd] = None):
"""Creates MetadataWriter based on general/input/output information.
Args:
model_buffer: valid buffer of the model file.
general_md: general information about the model. If not specified, default
general metadata will be generated.
input_md: input text tensor information, if not specified, default input
metadata will be generated.
output_md: output classification tensor information, if not specified,
default output metadata will be generated.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputTextTensorMd(
name=_INPUT_NAME, description=_INPUT_DESCRIPTION)
if output_md is None:
output_md = metadata_info.ClassificationTensorMd(
name=_OUTPUT_NAME, description=_OUTPUT_DESCRIPTION)
if output_md.associated_files is None:
output_md.associated_files = []
tokenizer_files = []
if input_md.tokenizer_md:
tokenizer_files = writer_utils.get_tokenizer_associated_files(
input_md.tokenizer_md.create_metadata().options)
return super().create_from_metadata_info(
model_buffer=model_buffer,
general_md=general_md,
input_md=[input_md],
output_md=[output_md],
associated_files=[
file.file_path for file in output_md.associated_files
] + tokenizer_files)
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputImageTensorMd] = None,
output_md: Optional[metadata_info.TensorMd] = None):
"""Creates MetadataWriter based on general/input/outputs information.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model.
input_md: input image tensor informaton.
output_md: output segmentation mask tensor informaton. This tensor is a
multidimensional array of [1 x mask_height x mask_width x num_classes],
where mask_width and mask_height are the dimensions of the segmentation
masks produced by the model, and num_classes is the number of classes
supported by the model.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputImageTensorMd(
name=_INPUT_NAME,
description=_INPUT_DESCRIPTION,
color_space_type=_metadata_fb.ColorSpaceType.RGB)
if output_md is None:
output_md = metadata_info.TensorMd(name=_OUTPUT_NAME,
description=_OUTPUT_DESCRIPTION)
if output_md.associated_files is None:
output_md.associated_files = []
return super().create_from_metadata(
model_buffer,
model_metadata=general_md.create_metadata(),
input_metadata=[input_md.create_metadata()],
output_metadata=[_create_segmentation_masks_metadata(output_md)],
associated_files=[
file.file_path for file in output_md.associated_files
])
def create_from_metadata_info_for_multihead(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputAudioTensorMd] = None,
output_md_list: Optional[List[
metadata_info.ClassificationTensorMd]] = None):
"""Creates a MetadataWriter instance for multihead models.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model. If not specified, default
general metadata will be generated.
input_md: input audio tensor informaton. If not specified, default input
metadata will be generated.
output_md_list: information of each output tensor head. If not specified,
default metadata will be generated for each output tensor. If
`tensor_name` in each `ClassificationTensorMd` instance is not
specified, elements in `output_md_list` need to have one-to-one mapping
with the output tensors [1] in the TFLite model.
[1]:
https://github.com/tensorflow/tflite-support/blob/b2a509716a2d71dfff706468680a729cc1604cff/tensorflow_lite_support/metadata/metadata_schema.fbs#L605-L612
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputAudioTensorMd(
name=_INPUT_NAME, description=_INPUT_DESCRIPTION)
associated_files = []
for md in output_md_list or []:
associated_files.extend(
[file.file_path for file in md.associated_files or []])
return super().create_from_metadata_info(
model_buffer=model_buffer,
general_md=general_md,
input_md=[input_md],
output_md=output_md_list,
associated_files=associated_files)
def test_create_metadata_should_succeed(self):
general_md = metadata_info.GeneralMd(name="model",
version="v1",
description="A ML model.",
author="TensorFlow",
licenses="Apache")
general_metadata = general_md.create_metadata()
# Create the Flatbuffers object and convert it to the json format.
builder = flatbuffers.Builder(0)
builder.Finish(general_metadata.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_json = _metadata.convert_to_json(bytes(builder.Output()))
expected_json = test_utils.load_file(self._EXPECTED_GENERAL_META_JSON,
"r")
self.assertEqual(metadata_json, expected_json)
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputImageTensorMd] = None,
output_md: Optional[metadata_info.ClassificationTensorMd] = None):
"""Creates MetadataWriter based on general/input/output information.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model. If not specified, default
general metadata will be generated.
input_md: input image tensor informaton, if not specified, default input
metadata will be generated.
output_md: output classification tensor informaton, if not specified,
default output metadata will be generated.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputImageTensorMd(
name=_INPUT_NAME,
description=_INPUT_DESCRIPTION,
color_space_type=_metadata_fb.ColorSpaceType.RGB)
if output_md is None:
output_md = metadata_info.ClassificationTensorMd(
name=_OUTPUT_NAME, description=_OUTPUT_DESCRIPTION)
if output_md.associated_files is None:
output_md.associated_files = []
return super().create_from_metadata_info(
model_buffer=model_buffer,
general_md=general_md,
input_md=[input_md],
output_md=[output_md],
associated_files=[
file.file_path for file in output_md.associated_files
])
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[List[Type[metadata_info.TensorMd]]] = None,
output_md: Optional[List[Type[metadata_info.TensorMd]]] = None,
associated_files: Optional[List[str]] = None):
"""Creates MetadataWriter based on the metadata information.
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.
associated_files: path to the associated files to be populated.
Returns:
A MetadataWriter Object.
Raises:
ValueError: if the tensor names from `input_md` and `output_md` do not
match the tensor names read from the model.
"""
if general_md is None:
general_md = metadata_info.GeneralMd()
if input_md is None:
input_md = []
if output_md is None:
output_md = []
# Order the input/output metadata according to tensor orders from the model.
input_md = _order_tensor_metadata(
input_md, writer_utils.get_input_tensor_names(model_buffer))
output_md = _order_tensor_metadata(
output_md, writer_utils.get_output_tensor_names(model_buffer))
model_metadata = general_md.create_metadata()
input_metadata = [m.create_metadata() for m in input_md]
output_metadata = [m.create_metadata() for m in output_md]
return cls.create_from_metadata(model_buffer, model_metadata,
input_metadata, output_metadata,
associated_files)
def test_create_from_metadata_info_succeeds_for_multihead(self):
calibration_file1 = test_utils.create_calibration_file(
self.get_temp_dir(), "score_cali_1.txt")
calibration_file2 = test_utils.create_calibration_file(
self.get_temp_dir(), "score_cali_2.txt")
general_md = metadata_info.GeneralMd(name="AudioClassifier")
input_md = metadata_info.InputAudioTensorMd(
name="audio_clip", sample_rate=_SAMPLE_RATE, channels=_CHANNELS)
# The output tensors in the model are: Identity, Identity_1
# Create metadata in a different order to test if MetadataWriter can correct
# it.
output_head_md_1 = metadata_info.ClassificationTensorMd(
name="head1",
label_files=[
metadata_info.LabelFileMd("labels_en_1.txt"),
metadata_info.LabelFileMd("labels_cn_1.txt")
],
score_calibration_md=metadata_info.ScoreCalibrationMd(
_metadata_fb.ScoreTransformationType.LOG,
_DEFAULT_SCORE_CALIBRATION_VALUE, calibration_file1),
tensor_name="Identity_1")
output_head_md_2 = metadata_info.ClassificationTensorMd(
name="head2",
label_files=[
metadata_info.LabelFileMd("labels_en_2.txt"),
metadata_info.LabelFileMd("labels_cn_2.txt")
],
score_calibration_md=metadata_info.ScoreCalibrationMd(
_metadata_fb.ScoreTransformationType.LOG,
_DEFAULT_SCORE_CALIBRATION_VALUE, calibration_file2),
tensor_name="Identity")
writer = (
audio_classifier.MetadataWriter.create_from_metadata_info_for_multihead(
test_utils.load_file(_MULTIHEAD_MODEL), general_md, input_md,
[output_head_md_1, output_head_md_2]))
metadata_json = writer.get_metadata_json()
expected_json = test_utils.load_file(_JSON_MULTIHEAD, "r")
self.assertEqual(metadata_json, expected_json)
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputImageTensorMd] = None,
output_location_md: Optional[metadata_info.TensorMd] = None,
output_category_md: Optional[
metadata_info.CategoryTensorMd] = None,
output_score_md: Union[
None, metadata_info.TensorMd,
metadata_info.ClassificationTensorMd] = None,
output_number_md: Optional[metadata_info.TensorMd] = None):
"""Creates MetadataWriter based on general/input/outputs information.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model.
input_md: input image tensor informaton.
output_location_md: output location tensor informaton. The location tensor
is a multidimensional array of [N][4] floating point values between 0
and 1, the inner arrays representing bounding boxes in the form [top,
left, bottom, right].
output_category_md: output category tensor information. The category
tensor is an array of N integers (output as floating point values) each
indicating the index of a class label from the labels file.
output_score_md: output score tensor information. The score tensor is an
array of N floating point values between 0 and 1 representing
probability that a class was detected. Use ClassificationTensorMd to
calibrate score.
output_number_md: output number of detections tensor information. This
tensor is an integer value of N.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputImageTensorMd(
name=_INPUT_NAME,
description=_INPUT_DESCRIPTION,
color_space_type=_metadata_fb.ColorSpaceType.RGB)
warn_message_format = (
"The output name isn't the default string \"%s\". This may cause the "
"model not work in the TFLite Task Library since the tensor name will "
"be used to handle the output order in the TFLite Task Library.")
if output_location_md is None:
output_location_md = metadata_info.TensorMd(
name=_OUTPUT_LOCATION_NAME,
description=_OUTPUT_LOCATION_DESCRIPTION)
elif output_location_md.name != _OUTPUT_LOCATION_NAME:
logging.warning(warn_message_format, _OUTPUT_LOCATION_NAME)
if output_category_md is None:
output_category_md = metadata_info.CategoryTensorMd(
name=_OUTPUT_CATRGORY_NAME,
description=_OUTPUT_CATEGORY_DESCRIPTION)
elif output_category_md.name != _OUTPUT_CATRGORY_NAME:
logging.warning(warn_message_format, _OUTPUT_CATRGORY_NAME)
if output_score_md is None:
output_score_md = metadata_info.ClassificationTensorMd(
name=_OUTPUT_SCORE_NAME,
description=_OUTPUT_SCORE_DESCRIPTION,
)
elif output_score_md.name != _OUTPUT_SCORE_NAME:
logging.warning(warn_message_format, _OUTPUT_SCORE_NAME)
if output_number_md is None:
output_number_md = metadata_info.TensorMd(
name=_OUTPUT_NUMBER_NAME,
description=_OUTPUT_NUMBER_DESCRIPTION)
elif output_number_md.name != _OUTPUT_NUMBER_NAME:
logging.warning(warn_message_format, _OUTPUT_NUMBER_NAME)
# Create output tensor group info.
group = _metadata_fb.TensorGroupT()
group.name = _GROUP_NAME
group.tensorNames = [
output_location_md.name, output_category_md.name,
output_score_md.name
]
# Gets the tensor inidces of tflite outputs and then gets the order of the
# output metadata by the value of tensor indices. For instance, if the
# output indices are [601, 599, 598, 600], tensor names and indices aligned
# are:
# - location: 598
# - category: 599
# - score: 600
# - number of detections: 601
# because of the op's ports of TFLITE_DETECTION_POST_PROCESS
# (https://github.com/tensorflow/tensorflow/blob/a4fe268ea084e7d323133ed7b986e0ae259a2bc7/tensorflow/lite/kernels/detection_postprocess.cc#L47-L50).
# Thus, the metadata of tensors are sorted in this way, according to
# output_tensor_indicies correctly.
output_tensor_indices = _get_tflite_outputs(model_buffer)
metadata_list = [
_create_location_metadata(output_location_md),
_create_metadata_with_value_range(output_category_md),
_create_metadata_with_value_range(output_score_md),
output_number_md.create_metadata()
]
# Align indices with tensors.
sorted_indices = sorted(output_tensor_indices)
indices_to_tensors = dict(zip(sorted_indices, metadata_list))
# Output metadata according to output_tensor_indices.
output_metadata = [
indices_to_tensors[i] for i in output_tensor_indices
]
# Create subgraph info.
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = [input_md.create_metadata()]
subgraph_metadata.outputTensorMetadata = output_metadata
subgraph_metadata.outputTensorGroups = [group]
# Create model metadata
model_metadata = general_md.create_metadata()
model_metadata.subgraphMetadata = [subgraph_metadata]
b = flatbuffers.Builder(0)
b.Finish(model_metadata.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
associated_files = []
_extend_new_files(associated_files,
output_category_md.associated_files)
_extend_new_files(associated_files, output_score_md.associated_files)
return cls(model_buffer, b.Output(), associated_files=associated_files)
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputImageTensorMd] = None,
output_location_md: Optional[metadata_info.TensorMd] = None,
output_category_md: Optional[metadata_info.CategoryTensorMd] = None,
output_score_md: Optional[metadata_info.TensorMd] = None,
output_number_md: Optional[metadata_info.TensorMd] = None):
"""Creates MetadataWriter based on general/input/outputs information.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model.
input_md: input image tensor informaton.
output_location_md: output location tensor informaton. The location tensor
is a multidimensional array of [N][4] floating point values between 0
and 1, the inner arrays representing bounding boxes in the form [top,
left, bottom, right].
output_category_md: output category tensor information. The category
tensor is an array of N integers (output as floating point values) each
indicating the index of a class label from the labels file.
output_score_md: output score tensor information. The score tensor is an
array of N floating point values between 0 and 1 representing
probability that a class was detected.
output_number_md: output number of dections tensor information. This
tensor is an integer value of N.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputImageTensorMd(
name=_INPUT_NAME,
description=_INPUT_DESCRIPTION,
color_space_type=_metadata_fb.ColorSpaceType.RGB)
if output_location_md is None:
output_location_md = metadata_info.TensorMd(
name=_OUTPUT_LOCATION_NAME, description=_OUTPUT_LOCATION_DESCRIPTION)
if output_category_md is None:
output_category_md = metadata_info.CategoryTensorMd(
name=_OUTPUT_CATRGORY_NAME, description=_OUTPUT_CATEGORY_DESCRIPTION)
if output_score_md is None:
output_score_md = metadata_info.TensorMd(
name=_OUTPUT_SCORE_NAME, description=_OUTPUT_SCORE_DESCRIPTION)
if output_number_md is None:
output_number_md = metadata_info.TensorMd(
name=_OUTPUT_NUMBER_NAME, description=_OUTPUT_NUMBER_DESCRIPTION)
if output_category_md.associated_files is None:
output_category_md.associated_files = []
return super().create_from_metadata(
model_buffer,
model_metadata=general_md.create_metadata(),
input_metadata=[input_md.create_metadata()],
output_metadata=[
_create_location_metadata(output_location_md),
_create_metadata_with_value_range(output_category_md),
_create_metadata_with_value_range(output_score_md),
output_number_md.create_metadata()
],
associated_files=[
file.file_path for file in output_category_md.associated_files
])
def create_from_metadata_info(
cls,
model_buffer: bytearray,
general_md: Optional[metadata_info.GeneralMd] = None,
input_md: Optional[metadata_info.InputImageTensorMd] = None,
output_location_md: Optional[metadata_info.TensorMd] = None,
output_category_md: Optional[metadata_info.CategoryTensorMd] = None,
output_score_md: Optional[metadata_info.TensorMd] = None,
output_number_md: Optional[metadata_info.TensorMd] = None):
"""Creates MetadataWriter based on general/input/outputs information.
Args:
model_buffer: valid buffer of the model file.
general_md: general infromation about the model.
input_md: input image tensor informaton.
output_location_md: output location tensor informaton. The location tensor
is a multidimensional array of [N][4] floating point values between 0
and 1, the inner arrays representing bounding boxes in the form [top,
left, bottom, right].
output_category_md: output category tensor information. The category
tensor is an array of N integers (output as floating point values) each
indicating the index of a class label from the labels file.
output_score_md: output score tensor information. The score tensor is an
array of N floating point values between 0 and 1 representing
probability that a class was detected.
output_number_md: output number of dections tensor information. This
tensor is an integer value of N.
Returns:
A MetadataWriter object.
"""
if general_md is None:
general_md = metadata_info.GeneralMd(
name=_MODEL_NAME, description=_MODEL_DESCRIPTION)
if input_md is None:
input_md = metadata_info.InputImageTensorMd(
name=_INPUT_NAME,
description=_INPUT_DESCRIPTION,
color_space_type=_metadata_fb.ColorSpaceType.RGB)
if output_location_md is None:
output_location_md = metadata_info.TensorMd(
name=_OUTPUT_LOCATION_NAME, description=_OUTPUT_LOCATION_DESCRIPTION)
if output_category_md is None:
output_category_md = metadata_info.CategoryTensorMd(
name=_OUTPUT_CATRGORY_NAME, description=_OUTPUT_CATEGORY_DESCRIPTION)
if output_score_md is None:
output_score_md = metadata_info.TensorMd(
name=_OUTPUT_SCORE_NAME, description=_OUTPUT_SCORE_DESCRIPTION)
if output_number_md is None:
output_number_md = metadata_info.TensorMd(
name=_OUTPUT_NUMBER_NAME, description=_OUTPUT_NUMBER_DESCRIPTION)
if output_category_md.associated_files is None:
output_category_md.associated_files = []
# Create output tensor group info.
group = _metadata_fb.TensorGroupT()
group.name = _GROUP_NAME
group.tensorNames = [
output_location_md.name, output_category_md.name, output_score_md.name
]
# Create subgraph info.
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = [input_md.create_metadata()]
subgraph_metadata.outputTensorMetadata = [
_create_location_metadata(output_location_md),
_create_metadata_with_value_range(output_category_md),
_create_metadata_with_value_range(output_score_md),
output_number_md.create_metadata()
]
subgraph_metadata.outputTensorGroups = [group]
# Create model metadata
model_metadata = general_md.create_metadata()
model_metadata.subgraphMetadata = [subgraph_metadata]
b = flatbuffers.Builder(0)
b.Finish(
model_metadata.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return cls(
model_buffer,
b.Output(),
associated_files=[
file.file_path for file in output_category_md.associated_files
])
