"""Helper class to populate Audio Metadata, to be used in `with` statement.

Simple usage for model with two classification heads.

with MetadataWriter(tflite_path) as writer:

writer.add_input(sample_rate=16000, channels=1)

writer.add_output(name='animal_sound', labels=['dog', 'cat'])

writer.add_output(name='speech_command', labels=['yes', 'no'])

writer.save(tflite_path, json_filepath)

`add_output` can also take an ordered dict for multiple locales, example:

writer.add_output(name='animal_sound', labels=collections.OrderedDict([

('en', ['bird', 'cat']),

('fr', ['oiseau', 'chat'])

]))

"""

def __init__(self, tflite_filepath, **kwargs):

self._model = writer_utils.load_file(tflite_filepath)

self._general_md = md_info.GeneralMd(**kwargs)

self._inputs = []

self._outputs = []

def __enter__(self):

self._temp_folder = tempfile.TemporaryDirectory()

return self

def __exit__(self, exc_type, exc_val, exc_tb):

self._temp_folder.cleanup()

# Delete the attribute so that it errors out if not in `with` statement.

delattr(self, '_temp_folder')

def add_input(self, **kwargs):

"""Add metadta for the input tensor."""

self._inputs.append(md_info.InputAudioTensorMd(**kwargs))

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 if locale else ''))

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 save(self, tflite_filepath=None, json_filepath=None):

"""Persist model with metadata."""

if len(self._inputs) > 1:

raise ValueError('Only supports single input, got {}'.format(

len(self._inputs)))

input_md = self._inputs[0]

writer = md_writer.MetadataWriter.create_from_metadata_info_for_multihead(

model_buffer=self._model,

general_md=self._general_md,

input_md=input_md,

output_md_list=self._outputs)

if tflite_filepath:

writer_utils.save_file(writer.populate(), tflite_filepath, mode='wb')

if json_filepath:

writer_utils.save_file(

if version.parse(tf_version) < version.parse('2.5.0rc0'):

raise RuntimeError(

'Audio Tasks requires TF2.5 or later. For example, you can run the '

'following command to install TF2.5.0rc2:\n\n'

"""Base model spec for audio classification."""

def __init__(self, model_dir=None, strategy=None):

_ensure_tf25(_get_tf_version())

self.model_dir = model_dir

if not model_dir:

self.model_dir = tempfile.mkdtemp()

tf.compat.v1.logging.info('Checkpoints are stored in %s', self.model_dir)

self.strategy = strategy or tf.distribute.get_strategy()

@abc.abstractproperty

def target_sample_rate(self):

pass

@abc.abstractmethod

def create_model(self, num_classes, train_whole_model=False):

pass

@abc.abstractmethod

def run_classifier(self, model, epochs, train_ds, validation_ds, **kwargs):

pass

def preprocess_ds(self, ds, is_training=False, cache_fn=None):

"""Returns a preprocessed dataset."""

_ = is_training

_ = cache_fn

return ds

def get_default_quantization_config(self):

"""Gets the default quantization configuration."""

"""Load a model replicating WebAudio's AnalyzerNode.getFloatFrequencyData."""

model_name = 'sc_preproc_model'

file_extension = '.tar.gz'

filename = model_name + file_extension

# Load the preprocessing model, which transforms audio waveform into

# spectrograms (2D image-like representation of sound).

# This model replicates WebAudio's AnalyzerNode.getFloatFrequencyData

# (https://developer.mozilla.org/en-US/docs/Web/API/AnalyserNode/getFloatFrequencyData).

# It performs short-time Fourier transform (STFT) using a length-2048 Blackman

# window. It opeartes on mono audio at the 44100-Hz sample rate.

filepath = tf.keras.utils.get_file(

filename,

f'{TFJS_MODEL_ROOT}/speech-commands/conversion/{filename}',

cache_subdir='model_maker',

extract=True)

model_path = _remove_suffix_if_possible(filepath, file_extension)

"""Download TFJS speech command model for fine-tune."""

origin_root = f'{TFJS_MODEL_ROOT}/speech-commands/v0.3/browser_fft/18w'

files_to_download = [

'metadata.json', 'model.json', 'group1-shard1of2', 'group1-shard2of2'

]

for filename in files_to_download:

filepath = tf.keras.utils.get_file(

filename,

f'{origin_root}/{filename}',

cache_subdir='model_maker/tfjs-sc-model')

model_path = os.path.join(os.path.dirname(filepath), 'model.json')

"""Model good at detecting speech commands, using Browser FFT spectrum."""

EXPECTED_WAVEFORM_LENGTH = 44032

# Information used to populate TFLite metadata.

_MODEL_NAME = 'AudioClassifier'

_MODEL_DESCRIPTION = ('Identify the most prominent type in the audio clip '

'from a known set of categories.')

_MODEL_VERSION = 'v1'

_MODEL_AUTHOR = 'TensorFlow Lite Model Maker'

_MODEL_LICENSES = ('Apache License. Version 2.0 '

'http://www.apache.org/licenses/LICENSE-2.0.')

_SAMPLE_RATE = 44100

_CHANNELS = 1

_INPUT_NAME = 'audio_clip'

_INPUT_DESCRIPTION = 'Input audio clip to be classified.'

_OUTPUT_NAME = 'probability'

_OUTPUT_DESCRIPTION = 'Scores of the labels respectively.'

def __init__(self, model_dir=None, strategy=None):

"""Initialize a new instance for BrowserFFT spec.

Args:

model_dir: The location to save the model checkpoint files.

strategy: An instance of TF distribute strategy. If none, it will use the

default strategy (either SingleDeviceStrategy or the current scoped

strategy.

"""

super(BrowserFFTSpec, self).__init__(model_dir, strategy)

self._preprocess_model = _load_browser_fft_preprocess_model()

self._tfjs_sc_model = _load_tfjs_speech_command_model()

@property

def target_sample_rate(self):

return 44100

@tf.function(input_signature=[

tf.TensorSpec(shape=[None], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _ensure_length(self, wav, unused_label):

return len(wav) >= self.EXPECTED_WAVEFORM_LENGTH

@tf.function(input_signature=[

tf.TensorSpec(shape=[None], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _split(self, wav, label):

"""Split the long audio samples into multiple trunks."""

# wav shape: (audio_samples, )

chunks = tf.math.floordiv(len(wav), self.EXPECTED_WAVEFORM_LENGTH)

unused = tf.math.floormod(len(wav), self.EXPECTED_WAVEFORM_LENGTH)

# Drop unused data

wav = wav[:len(wav) - unused]

# Split the audio sample into multiple chunks

wav = tf.reshape(wav, (chunks, 1, self.EXPECTED_WAVEFORM_LENGTH))

return wav, tf.repeat(tf.expand_dims(label, 0), len(wav))

@tf.function(input_signature=[

tf.TensorSpec(shape=[1, EXPECTED_WAVEFORM_LENGTH], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _preprocess(self, x, label):

"""Preprocess the dataset to extract the spectrum."""

# Add small Gaussian noise to the input x

# to solve the potential "nan" problem of the preprocess_model.

x = x + 1e-05 * tf.random.normal(x.shape)

# x has shape (1, EXPECTED_WAVEFORM_LENGTH)

spectrum = self._preprocess_model(x)

# y has shape (1, embedding_len)

spectrum = tf.squeeze(spectrum, axis=0)

# y has shape (embedding_len,)

return spectrum, label

def preprocess_ds(self, ds, is_training=False, cache_fn=None):

del is_training

autotune = tf.data.AUTOTUNE

ds = ds.filter(self._ensure_length)

ds = ds.map(self._split, num_parallel_calls=autotune).unbatch()

ds = ds.map(self._preprocess, num_parallel_calls=autotune)

if cache_fn:

ds = cache_fn(ds)

return ds

def create_model(self, num_classes, train_whole_model=False):

if num_classes <= 1:

raise ValueError(

'AudioClassifier expects `num_classes` to be greater than 1')

model = tf.keras.Sequential()

for layer in self._tfjs_sc_model.layers[:-1]:

model.add(layer)

model.add(

tf.keras.layers.Dense(

name='classification_head', units=num_classes,

activation='softmax'))

if not train_whole_model:

# Freeze all but the last layer of the model. The last layer will be

# fine-tuned during transfer learning.

for layer in model.layers[:-1]:

layer.trainable = False

return model

def run_classifier(self, model, epochs, train_ds, validation_ds, **kwargs):

model.compile(

optimizer='adam', loss='categorical_crossentropy', metrics=['acc'])

hist = model.fit(

train_ds, validation_data=validation_ds, epochs=epochs, **kwargs)

return hist

def create_serving_model(self, training_model):

"""Create a model for serving."""

combined = tf.keras.Sequential()

combined.add(self._preprocess_model)

combined.add(training_model)

# Build the model.

combined.build([None, self.EXPECTED_WAVEFORM_LENGTH])

return combined

def _export_metadata(self, tflite_filepath, index_to_label,

export_metadata_json_file):

"""Export TFLite metadata."""

with MetadataWriter(

tflite_filepath,

name=self._MODEL_NAME,

description=self._MODEL_DESCRIPTION,

version=self._MODEL_VERSION,

author=self._MODEL_AUTHOR,

licenses=self._MODEL_LICENSES) as writer:

writer.add_input(

name=self._INPUT_NAME,

description=self._INPUT_DESCRIPTION,

sample_rate=self._SAMPLE_RATE,

channels=self._CHANNELS)

writer.add_output(

labels=index_to_label,

name=self._OUTPUT_NAME,

description=self._OUTPUT_DESCRIPTION)

json_filepath = (os.path.splitext(tflite_filepath)[0] +

'.json') if export_metadata_json_file else None

writer.save(tflite_filepath, json_filepath)

def export_tflite(self,

model,

tflite_filepath,

with_metadata=True,

export_metadata_json_file=True,

index_to_label=None,

quantization_config=None):

"""Converts the retrained model to tflite format and saves it.

This method overrides the default `CustomModel._export_tflite` method, and

include the pre-processing in the exported TFLite library since support

library can't handle audio tasks yet.

Args:

model: An instance of the keras classification model to be exported.

tflite_filepath: File path to save tflite model.

with_metadata: Whether the output tflite model contains metadata.

export_metadata_json_file: Whether to export metadata in json file. If

True, export the metadata in the same directory as tflite model.Used

only if `with_metadata` is True.

index_to_label: A list that map from index to label class name.

quantization_config: Configuration for post-training quantization.

"""

combined = self.create_serving_model(model)

# Sets batch size from None to 1 when converting to tflite.

model_util.set_batch_size(model, batch_size=1)

model_util.export_tflite(

combined, tflite_filepath, quantization_config=quantization_config)

# Sets batch size back to None to support retraining later.

model_util.set_batch_size(model, batch_size=None)

if with_metadata:

if not ENABLE_METADATA:

print('Writing Metadata is not support in the installed tflite-support '

'version. Please use tflite-support >= 0.2.*')

else:

self._export_metadata(tflite_filepath, index_to_label,

"""Model good at detecting environmental sounds, using YAMNet embedding."""

EXPECTED_WAVEFORM_LENGTH = 15600 # effectively 0.975s

EMBEDDING_SIZE = 1024

# Information used to populate TFLite metadata.

_MODEL_NAME = 'yamnet/classification'

_MODEL_DESCRIPTION = 'Recognizes sound events'

_MODEL_VERSION = 'v1'

_MODEL_AUTHOR = 'TensorFlow Lite Model Maker'

_MODEL_LICENSES = ('Apache License. Version 2.0 '

'http://www.apache.org/licenses/LICENSE-2.0.')

_SAMPLE_RATE = 16000

_CHANNELS = 1

_INPUT_NAME = 'audio_clip'

_INPUT_DESCRIPTION = 'Input audio clip to be classified.'

_YAMNET_OUTPUT_NAME = 'yamnet'

_YAMNET_OUTPUT_DESCRIPTION = ('Scores in range 0..1.0 for each of the 521 '

'output classes.')

_CUSTOM_OUTPUT_NAME = 'custom'

_CUSTOM_OUTPUT_DESCRIPTION = (

'Scores in range 0..1.0 for each output classes.')

def __init__(

self,

model_dir: None = None,

strategy: None = None,

yamnet_model_handle='https://tfhub.dev/google/yamnet/1',

frame_length=EXPECTED_WAVEFORM_LENGTH, # Window size 0.975 s

frame_step=EXPECTED_WAVEFORM_LENGTH // 2, # Hop of 0.975 /2 s

keep_yamnet_and_custom_heads=True):

"""Initialize a new instance for YAMNet spec.

Args:

model_dir: The location to save the model checkpoint files.

strategy: An instance of TF distribute strategy. If none, it will use the

default strategy (either SingleDeviceStrategy or the current scoped

strategy.

yamnet_model_handle: Path of the TFHub model for retrining.

frame_length: The number of samples in each audio frame. If the audio file

is shorter than `frame_length`, then the audio file will be ignored.

frame_step: The number of samples between two audio frames. This value

should be smaller than `frame_length`, otherwise some samples will be

ignored.

keep_yamnet_and_custom_heads: Boolean, decides if the final TFLite model

contains both YAMNet and custom trained classification heads. When set

to False, only the trained custom head will be preserved.

"""

super(YAMNetSpec, self).__init__(model_dir, strategy)

self._yamnet_model_handle = yamnet_model_handle

self._yamnet_model = hub.load(yamnet_model_handle)

self._frame_length = frame_length

self._frame_step = frame_step

self._keep_yamnet_and_custom_heads = keep_yamnet_and_custom_heads

@property

def target_sample_rate(self):

return self._SAMPLE_RATE

def create_model(self, num_classes, train_whole_model=False):

model = tf.keras.Sequential([

tf.keras.layers.InputLayer(

input_shape=(YAMNetSpec.EMBEDDING_SIZE),

dtype=tf.float32,

name='embedding'),

tf.keras.layers.Dense(

num_classes, name='classification_head', activation='softmax')

])

return model

def run_classifier(self, model, epochs, train_ds, validation_ds, **kwargs):

model.compile(

optimizer='adam', loss='categorical_crossentropy', metrics=['acc'])

hist = model.fit(

train_ds, validation_data=validation_ds, epochs=epochs, **kwargs)

return hist

# Annotate the TF function with input_signature to avoid re-tracing. Otherwise

# the TF function gets retraced everytime the input shape is changed.

# Check https://www.tensorflow.org/api_docs/python/tf/function#args_1 for more

# information.

@tf.function(input_signature=[

tf.TensorSpec(shape=[None], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _frame(self, wav, label):

clips = tf.signal.frame(

wav, frame_length=self._frame_length, frame_step=self._frame_step)

batch_labels = tf.repeat(tf.expand_dims(label, 0), len(clips))

return clips, batch_labels

@tf.function(input_signature=[

tf.TensorSpec(shape=[None], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _extract_embedding(self, wav, label):

_, embeddings, _ = self._yamnet_model(wav) # (chunks, EMBEDDING_SIZE)

embedding = tf.reduce_mean(embeddings, axis=0)

return embedding, label

@tf.function(input_signature=[

tf.TensorSpec(shape=[EMBEDDING_SIZE], dtype=tf.float32),

tf.TensorSpec([], dtype=tf.int32)

])

def _add_noise(self, embedding, label):

noise = tf.random.normal(

embedding.shape, mean=0.0, stddev=.2, dtype=tf.dtypes.float32)

return noise + embedding, label

def preprocess_ds(self, ds, is_training=False, cache_fn=None):

autotune = tf.data.AUTOTUNE

ds = ds.map(self._frame, num_parallel_calls=autotune).unbatch()

ds = ds.map(self._extract_embedding, num_parallel_calls=autotune)

# Cache intermediate results right before data augmentation.

if cache_fn:

ds = cache_fn(ds)

if is_training:

ds = ds.map(self._add_noise, num_parallel_calls=autotune)

return ds

def _yamnet_labels(self):

class_map_path = self._yamnet_model.class_map_path().numpy()

class_map_csv_text = tf.io.read_file(class_map_path).numpy().decode('utf-8')

class_map_csv = io.StringIO(class_map_csv_text)

class_names = [

display_name for (class_index, mid,

display_name) in csv.reader(class_map_csv)

]

class_names = class_names[1:] # Skip CSV header

return class_names

def _export_metadata(self, tflite_filepath, index_to_label,

export_metadata_json_file):

"""Export TFLite metadata."""

with MetadataWriter(

tflite_filepath,

name=self._MODEL_NAME,

description=self._MODEL_DESCRIPTION,

version=self._MODEL_VERSION,

author=self._MODEL_AUTHOR,

licenses=self._MODEL_LICENSES) as writer:

writer.add_input(

name=self._INPUT_NAME,

description=self._INPUT_DESCRIPTION,

sample_rate=self._SAMPLE_RATE,

channels=self._CHANNELS)

if self._keep_yamnet_and_custom_heads:

writer.add_output(

labels=self._yamnet_labels(),

name=self._YAMNET_OUTPUT_NAME,

description=self._YAMNET_OUTPUT_DESCRIPTION)

writer.add_output(

labels=index_to_label,

name=self._CUSTOM_OUTPUT_NAME,

description=self._CUSTOM_OUTPUT_DESCRIPTION)

json_filepath = (os.path.splitext(tflite_filepath)[0] +

'.json') if export_metadata_json_file else None

writer.save(tflite_filepath, json_filepath)

def create_serving_model(self, training_model):

"""Create a model for serving."""

embedding_extraction_layer = hub.KerasLayer(

self._yamnet_model_handle, trainable=False)

keras_input = tf.keras.Input(

shape=(YAMNetSpec.EXPECTED_WAVEFORM_LENGTH,),

dtype=tf.float32,

name='audio') # (1, wav)

reshaped_input = tf.reshape(keras_input,

(YAMNetSpec.EXPECTED_WAVEFORM_LENGTH,)) # (wav)

scores, embeddings, _ = embedding_extraction_layer(reshaped_input)

serving_outputs = training_model(embeddings)

if self._keep_yamnet_and_custom_heads:

serving_model = tf.keras.Model(keras_input, [scores, serving_outputs])

else:

serving_model = tf.keras.Model(keras_input, serving_outputs)

return serving_model

def export_tflite(self,

model,

tflite_filepath,

with_metadata=True,

export_metadata_json_file=True,

index_to_label=None,

quantization_config=None):

"""Converts the retrained model to tflite format and saves it.

This method overrides the default `CustomModel._export_tflite` method, and

include the spectrom extraction in the model.

The exported model has input shape (1, number of wav samples)

Args:

model: An instance of the keras classification model to be exported.

tflite_filepath: File path to save tflite model.

with_metadata: Whether the output tflite model contains metadata.

export_metadata_json_file: Whether to export metadata in json file. If

True, export the metadata in the same directory as tflite model. Used

only if `with_metadata` is True.

index_to_label: A list that map from index to label class name.

quantization_config: Configuration for post-training quantization.

"""

serving_model = self.create_serving_model(model)

# TODO(b/164229433): Remove SELECT_TF_OPS once changes in the bug are

# released.

model_util.export_tflite(

serving_model, tflite_filepath, quantization_config=quantization_config)

if with_metadata:

if not ENABLE_METADATA:

print('Writing Metadata is not support in the current tflite-support '

'version. Please use tflite-support >= 0.2.*')

else:

self._export_metadata(tflite_filepath, index_to_label,