def add_output(self, name, labels, **kwargs):
"""Add metadata for output tensor in order."""
if isinstance(labels, list):
default_locale = None
labels = collections.OrderedDict([(default_locale, labels)])
return self.add_output(name, labels, **kwargs)
label_files = []
if isinstance(labels, collections.OrderedDict):
for locale, label_list in labels.items():
full_path = os.path.join(
self._temp_folder.name,
'{}_labels_{}.txt'.format(name, locale or 'default'))
model_util.export_labels(full_path, label_list)
label_files.append(
md_info.LabelFileMd(file_path=full_path, locale=locale))
else:
raise ValueError(
'`labels` should be either a list of labels or an ordered dict mapping `locale` -> list of labels. got: {}'
.format(labels))
idx = len(self._outputs)
self._outputs.append(
md_info.ClassificationTensorMd(
name=name,
label_files=label_files,
tensor_type=writer_utils.get_output_tensor_types(
self._model)[idx],
**kwargs))
def _export_metadata(self, tflite_filepath, index_to_label,
export_metadata_json_file):
"""Export TFLite metadata."""
with tempfile.TemporaryDirectory() as temp_dir:
# Prepare metadata
with open(tflite_filepath, 'rb') as f:
model_buffer = f.read()
general_md = md_info.GeneralMd(name=self._MODEL_NAME,
description=self._MODEL_DESCRIPTION,
version=self._MODEL_VERSION,
author=self._MODEL_AUTHOR,
licenses=self._MODEL_LICENSES)
input_md = md_info.InputAudioTensorMd(self._INPUT_NAME,
self._INPUT_DESCRIPTION,
self._SAMPLE_RATE,
self._CHANNELS)
# Save label files.
custom_label_filepath = os.path.join(temp_dir,
self._CUSTOM_LABEL_FILE)
self._export_labels(custom_label_filepath, index_to_label)
custom_output_md = md_info.ClassificationTensorMd(
name=self._CUSTOM_OUTPUT_NAME,
description=self._CUSTOM_OUTPUT_DESCRIPTION,
label_files=[
md_info.LabelFileMd(file_path=os.path.join(
temp_dir, self._CUSTOM_LABEL_FILE))
],
tensor_type=writer_utils.get_output_tensor_types(
model_buffer)[-1],
score_calibration_md=None)
if self._keep_yamnet_and_custom_heads:
yamnet_label_filepath = os.path.join(temp_dir,
self._YAMNET_LABEL_FILE)
self._export_labels(yamnet_label_filepath,
self._yamnet_labels())
yamnet_output_md = md_info.ClassificationTensorMd(
name=self._YAMNET_OUTPUT_NAME,
description=self._YAMNET_OUTPUT_DESCRIPTION,
label_files=[
md_info.LabelFileMd(file_path=os.path.join(
temp_dir, self._YAMNET_LABEL_FILE))
],
tensor_type=writer_utils.get_output_tensor_types(
model_buffer)[0],
score_calibration_md=None)
output_md = [yamnet_output_md, custom_output_md]
else:
output_md = [custom_output_md]
# Populate metadata
writer = md_writer.MetadataWriter.create_from_metadata_info_for_multihead(
model_buffer=model_buffer,
general_md=general_md,
input_md=input_md,
output_md_list=output_md)
output_model = writer.populate()
with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
f.write(output_model)
if export_metadata_json_file:
metadata_json = writer.get_metadata_json()
export_json_file = os.path.splitext(
tflite_filepath)[0] + '.json'
with open(export_json_file, 'w') as f:
f.write(metadata_json)