code: # Copyright 2020 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Audio model specification."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  import collections
  import csv
  import io
  import os
  import tempfile
  
  import tensorflow as tf
  from tensorflow_examples.lite.model_maker.core.api.api_util import mm_export
  from tensorflow_examples.lite.model_maker.core.task import model_util
  import tensorflow_hub as hub
  
  try:
    from tflite_support.metadata_writers import audio_classifier as md_writer  # pylint: disable=g-import-not-at-top
    from tflite_support.metadata_writers import metadata_info as md_info  # pylint: disable=g-import-not-at-top
    from tflite_support.metadata_writers import writer_utils  # pylint: disable=g-import-not-at-top
    ENABLE_METADATA = True
  except ImportError:
    ENABLE_METADATA = False
  
  
  class MetadataWriter:
    """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 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 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(
            writer.get_metadata_json(), json_filepath, mode='wt')
  
  
  def _ensure_tf25(version):
    if version < '2.5':
      raise RuntimeError(
          'Audio Tasks requires TF2.5 or later. For example, you can run the '
          'following command to install TF2.5.0rc2:\n\n'
          'pip3 install tensorflow==2.5.0rc2\n\n')
  
  
  def _get_tf_version():
    return tf.__version__
  
  
  class BaseSpec(abc.ABC):
    """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."""
      return None
  
  
  def _remove_suffix_if_possible(text, suffix):
    return text.rsplit(suffix, 1)[0]
  
  
  TFJS_MODEL_ROOT = 'https://storage.googleapis.com/tfjs-models/tfjs'
  
  
  def _load_browser_fft_preprocess_model():
    """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)
    return tf.keras.models.load_model(model_path)
  
  
  def _load_tfjs_speech_command_model():
    """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')
    return model_util.load_tfjs_keras_model(model_path)
  
  
  @mm_export('audio_classifier.BrowserFftSpec')
  class BrowserFFTSpec(BaseSpec):
    """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."""
      # 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):
      """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.
      """
      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=None)
  
      # 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,
                                export_metadata_json_file)
  
  
  @mm_export('audio_classifier.YamNetSpec')
  class YAMNetSpec(BaseSpec):
    """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 bigger than `frame_length`.
        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):
      """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.
      """
      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=None)
  
      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,
                                export_metadata_json_file)
  
apis: tensorflow.keras.Input tensorflow.reshape tensorflow.keras.models.load_model tensorflow.keras.utils.get_file tensorflow.keras.Sequential tensorflow.reduce_mean tensorflow.expand_dims tensorflow.keras.Model tensorflow.squeeze tensorflow.keras.layers.InputLayer tensorflow.random.normal tensorflow.keras.layers.Dense tensorflow.signal.frame tensorflow.TensorSpec tensorflow.distribute.get_strategy tensorflow.io.read_file tensorflow.compat.v1.logging.info