def _create_metadata_file(self):
associated_file1 = _metadata_fb.AssociatedFileT()
associated_file1.name = b"file1"
associated_file2 = _metadata_fb.AssociatedFileT()
associated_file2.name = b"file2"
self.expected_recorded_files = [
six.ensure_str(associated_file1.name),
six.ensure_str(associated_file2.name)
]
input_meta = _metadata_fb.TensorMetadataT()
output_meta = _metadata_fb.TensorMetadataT()
output_meta.associatedFiles = [associated_file2]
subgraph = _metadata_fb.SubGraphMetadataT()
# Create a model with two inputs and one output.
subgraph.inputTensorMetadata = [input_meta, input_meta]
subgraph.outputTensorMetadata = [output_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.name = "Mobilenet_quantized"
model_meta.associatedFiles = [associated_file1]
model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_file = self.create_tempfile().full_path
with open(metadata_file, "wb") as f:
f.write(b.Output())
return metadata_file
def testGetRecordedAssociatedFileListWithSubgraphProcessUnits(
self, tensor_type, tokenizer_type):
# Creates a metadata with the tokenizer in the subgraph process units.
tokenizer, expected_files = self._create_tokenizer(tokenizer_type)
# Create the subgraph with process units.
subgraph = _metadata_fb.SubGraphMetadataT()
if tensor_type is TensorType.INPUT:
subgraph.inputProcessUnits = [tokenizer]
elif tensor_type is TensorType.OUTPUT:
subgraph.outputProcessUnits = [tokenizer]
else:
raise ValueError(
"The tensor type, {0}, is unsupported.".format(tensor_type))
# Creates the input and output tensor meta to match self._model_file.
dummy_tensor_meta = _metadata_fb.TensorMetadataT()
subgraph.inputTensorMetadata = [dummy_tensor_meta, dummy_tensor_meta]
subgraph.outputTensorMetadata = [dummy_tensor_meta]
# Create a model metadata with the subgraph metadata
meta_buffer = self._create_model_meta_with_subgraph_meta(subgraph)
# Creates the tempfiles.
tempfiles = self._create_tempfiles(expected_files)
# Creates the MetadataPopulator object.
populator = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator.load_metadata_buffer(meta_buffer)
populator.load_associated_files(tempfiles)
populator.populate()
recorded_files = populator.get_recorded_associated_file_list()
self.assertEqual(set(recorded_files), set(expected_files))
def testPopulateMetadataFileToModelWithMetadataAndAssociatedFiles(self):
# First, creates a dummy metadata different from self._metadata_file. It
# needs to have the same input/output tensor numbers as self._model_file.
# Populates it and the associated files into the model.
input_meta = _metadata_fb.TensorMetadataT()
output_meta = _metadata_fb.TensorMetadataT()
subgraph = _metadata_fb.SubGraphMetadataT()
# Create a model with two inputs and one output.
subgraph.inputTensorMetadata = [input_meta, input_meta]
subgraph.outputTensorMetadata = [output_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.subgraphMetadata = [subgraph]
b = flatbuffers.Builder(0)
b.Finish(
model_meta.Pack(b),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
metadata_buf = b.Output()
# Populate the metadata.
populator1 = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator1.load_metadata_buffer(metadata_buf)
populator1.load_associated_files([self._file1, self._file2])
populator1.populate()
# Then, populate the metadata again.
populator2 = _metadata.MetadataPopulator.with_model_file(self._model_file)
populator2.load_metadata_file(self._metadata_file)
populator2.populate()
# Test if the metadata is populated correctly.
self._assert_golden_metadata(self._model_file)
def testPopulatedFullPathAssociatedFileShouldSucceed(self):
# Create AssociatedFileT using the full path file name.
associated_file = _metadata_fb.AssociatedFileT()
associated_file.name = self._file1
# Create model metadata with the associated file.
subgraph = _metadata_fb.SubGraphMetadataT()
subgraph.associatedFiles = [associated_file]
# Creates the input and output tensor metadata to match self._model_file.
dummy_tensor = _metadata_fb.TensorMetadataT()
subgraph.inputTensorMetadata = [dummy_tensor, dummy_tensor]
subgraph.outputTensorMetadata = [dummy_tensor]
md_buffer = self._create_model_meta_with_subgraph_meta(subgraph)
# Populate the metadata to a model.
populator = _metadata.MetadataPopulator.with_model_file(
self._model_file)
populator.load_metadata_buffer(md_buffer)
populator.load_associated_files([self._file1])
populator.populate()
# The recorded file name in metadata should only contain file basename; file
# directory should not be included.
recorded_files = populator.get_recorded_associated_file_list()
self.assertEqual(set(recorded_files),
set([os.path.basename(self._file1)]))
def create_from_metadata(
cls,
model_buffer: bytearray,
model_metadata: Optional[_metadata_fb.ModelMetadataT] = None,
input_metadata: Optional[List[
_metadata_fb.TensorMetadataT]] = None,
output_metadata: Optional[List[
_metadata_fb.TensorMetadataT]] = None,
associated_files: Optional[List[str]] = None):
"""Creates MetadataWriter based on the metadata Flatbuffers Python Objects.
Args:
model_buffer: valid buffer of the model file.
model_metadata: general model metadata [1]. The subgraph_metadata will be
refreshed with input_metadata and output_metadata.
input_metadata: a list of metadata of the input tensors [2].
output_metadata: a list of metadata of the output tensors [3].
associated_files: path to the associated files to be populated.
[1]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L640-L681
[2]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L590
[3]:
https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L599
Returns:
A MetadataWriter Object.
"""
# Create empty tensor metadata when input_metadata/output_metadata are None
# to bypass MetadataPopulator verification.
if not input_metadata:
model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
num_input_tensors = model.Subgraphs(0).InputsLength()
input_metadata = [_metadata_fb.TensorMetadataT()
] * num_input_tensors
if not output_metadata:
model = _schema_fb.Model.GetRootAsModel(model_buffer, 0)
num_output_tensors = model.Subgraphs(0).OutputsLength()
output_metadata = [_metadata_fb.TensorMetadataT()
] * num_output_tensors
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = input_metadata
subgraph_metadata.outputTensorMetadata = output_metadata
if model_metadata is None:
model_metadata = _metadata_fb.ModelMetadataT()
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)
def _create_dummy_model_metadata(
tensor_metadata: _metadata_fb.TensorMetadataT) -> bytes:
# Create a dummy model using the tensor metadata.
subgraph_metadata = _metadata_fb.SubGraphMetadataT()
subgraph_metadata.inputTensorMetadata = [tensor_metadata]
model_metadata = _metadata_fb.ModelMetadataT()
model_metadata.subgraphMetadata = [subgraph_metadata]
# Create the Flatbuffers object and convert it to the json format.
builder = flatbuffers.Builder(0)
builder.Finish(model_metadata.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
return bytes(builder.Output())
def testLoadMetadataBufferWithWrongOutputMetaNumberThrowsException(self):
# Create a dummy metadata with no output tensor metadata, while the expected
# number is 1.
input_meta = _metadata_fb.TensorMetadataT()
subgprah_meta = _metadata_fb.SubGraphMetadataT()
subgprah_meta.inputTensorMetadata = [input_meta, input_meta]
model_meta = _metadata_fb.ModelMetadataT()
model_meta.subgraphMetadata = [subgprah_meta]
builder = flatbuffers.Builder(0)
builder.Finish(
model_meta.Pack(builder),
_metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
meta_buf = builder.Output()
populator = _metadata.MetadataPopulator.with_model_buffer(self._model_buf)
with self.assertRaises(ValueError) as error:
populator.load_metadata_buffer(meta_buf)
self.assertEqual(
("The number of output tensors (1) should match the number of "
"output tensor metadata (0)"), str(error.exception))
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 = []
# 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
])
