def test_initialize_and_populate(self):
with mt.Writer(test_utils.load_file(_AUDIO_CLASSIFICATION_MODEL),
model_name='my_audio_model',
model_description='my_description') as writer:
out_dir = self.create_tempdir()
_, metadata_json = writer.populate(
os.path.join(out_dir, 'model.tflite'),
os.path.join(out_dir, 'metadata.json'))
self.assertJsonEqual(
metadata_json, """{
"name": "my_audio_model",
"description": "my_description",
"subgraph_metadata": [
{
"input_tensor_metadata": [
{
"name": "waveform_binary"
}
],
"output_tensor_metadata": [
{
"name": "tower0/network/layer32/final_output"
}
]
}
],
"min_parser_version": "1.0.0"
}
""")
def test_audio_embedder(self):
with mt.Writer(
test_utils.load_file(_AUDIO_EMBEDDING_MODEL),
model_name='audio_embedder',
model_description='Generate embedding for the input audio clip'
) as writer:
out_dir = self.create_tempdir()
writer.add_audio_input(sample_rate=16000, channels=1)
writer.add_embedding_output()
_, metadata_json = writer.populate(
os.path.join(out_dir, 'model.tflite'),
os.path.join(out_dir, 'metadata.json'))
self.assertEqual(
metadata_json, """{
"name": "audio_embedder",
"description": "Generate embedding for the input audio clip",
"subgraph_metadata": [
{
"input_tensor_metadata": [
{
"name": "audio",
"description": "Input audio clip to be processed.",
"content": {
"content_properties_type": "AudioProperties",
"content_properties": {
"sample_rate": 16000,
"channels": 1
}
},
"stats": {
}
}
],
"output_tensor_metadata": [
{
"name": "embedding",
"description": "Embedding vector of the input.",
"content": {
"content_properties_type": "FeatureProperties",
"content_properties": {
}
},
"stats": {
}
}
]
}
],
"min_parser_version": "1.3.0"
}
""")
def test_initialize_without_with_block(self):
writer = mt.Writer(test_utils.load_file(_AUDIO_CLASSIFICATION_MODEL),
model_name='test_model',
model_description='test_description')
# Calling `add_classification_output` outside the `with` block fails.
with self.assertRaisesRegex(AttributeError, '_temp_folder'):
writer.add_classification_output(mt.Labels().add(['cat', 'dog']))
writer.__enter__()
writer.add_classification_output(mt.Labels().add(['cat', 'dog']))
writer.__exit__(*sys.exc_info())
# Calling `add_classification_output` after `with` block closes also fails.
with self.assertRaisesRegex(AttributeError, '_temp_folder'):
writer.add_classification_output(mt.Labels().add(['cat', 'dog']))
def test_image_classifier(self):
with mt.Writer(
test_utils.load_file(_IMAGE_CLASSIFIER_MODEL),
model_name='image_classifier',
model_description='Imagenet classification model') as writer:
out_dir = self.create_tempdir()
writer.add_image_input(
norm_mean=[127.5, 127.5, 127.5],
norm_std=[127.5, 127.5, 127.5],
color_space_type=mt.Writer.color_space_types.RGB)
writer.add_classification_output(mt.Labels().add(['a', 'b', 'c']))
_, metadata_json = writer.populate(
os.path.join(out_dir, 'model.tflite'),
os.path.join(out_dir, 'metadat.json'))
self.assertEqual(
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_audio_classifier_with_locale_and_score_calibration(self):
with mt.Writer(
test_utils.load_file(_AUDIO_CLASSIFICATION_MODEL),
model_name='audio_classifier',
model_description='Classify the input audio clip') as writer:
out_dir = self.create_tempdir()
writer.add_audio_input(sample_rate=16000, channels=1)
writer.add_classification_output(
mt.Labels().add(['/id1', '/id2'],
use_as_category_name=True).add(
['sound1', 'sound2'],
'en').add(['son1', 'son2'], 'fr'),
score_calibration=mt.ScoreCalibration(
mt.ScoreCalibration.transformation_types.INVERSE_LOGISTIC,
[
mt.CalibrationParameter(1., 2., 3., None),
mt.CalibrationParameter(1., 2., 3., 4.),
],
default_score=0.5))
_, metadata_json = writer.populate(
os.path.join(out_dir, 'model.tflite'),
os.path.join(out_dir, 'metadata.tflite'))
self.assertEqual(
metadata_json, """{
"name": "audio_classifier",
"description": "Classify the input audio clip",
"subgraph_metadata": [
{
"input_tensor_metadata": [
{
"name": "audio",
"description": "Input audio clip to be processed.",
"content": {
"content_properties_type": "AudioProperties",
"content_properties": {
"sample_rate": 16000,
"channels": 1
}
},
"stats": {
}
}
],
"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.3.0"
}
""")
def test_audio_classifier_with_locale(self):
with mt.Writer(
test_utils.load_file(_AUDIO_CLASSIFICATION_MODEL),
model_name='audio_classifier',
model_description='Classify the input audio clip') as writer:
out_dir = self.create_tempdir()
writer.add_audio_input(sample_rate=16000, channels=1)
writer.add_classification_output(mt.Labels().add(
['/id1', '/id2'],
use_as_category_name=True).add(['sound1', 'sound2'],
'en').add(['son1', 'son2'],
'fr'))
_, metadata_json = writer.populate(
os.path.join(out_dir, 'model.tflite'),
os.path.join(out_dir, 'metadata.tflite'))
self.assertEqual(
metadata_json, """{
"name": "audio_classifier",
"description": "Classify the input audio clip",
"subgraph_metadata": [
{
"input_tensor_metadata": [
{
"name": "audio",
"description": "Input audio clip to be processed.",
"content": {
"content_properties_type": "AudioProperties",
"content_properties": {
"sample_rate": 16000,
"channels": 1
}
},
"stats": {
}
}
],
"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"
},
{
"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"
}
]
}
]
}
],
"min_parser_version": "1.3.0"
}
""")