tfrecords_write.py

import os
from typing import Union, Literal

import numpy as np
from colorama import Fore
from tensorflow import io as tf_io, train as tf_train

from pre_processing import LoadAudioFile, AudioAugmentation


class WriteTFRecord(LoadAudioFile, AudioAugmentation):
    def __init__(self,
                 sr: int = 16000,
                 start_time: float = 0.0,
                 duration=3, mono: bool = True,
                 pad: bool = True):
        LoadAudioFile.__init__(self, sr=sr,
                               start_time=start_time,
                               duration=duration,
                               mono=mono,
                               pad=pad)
        AudioAugmentation.__init__(self,)

        self.Extractor = None
        self.Normalizer = None

        self.feature_shape = None

    def create_tfrecord(self,
                        tfrecord_file_name: Union[str, os.PathLike],
                        path_datasets: list[np.ndarray],
                        names: list[Literal["real", "fake"]],
                        n: list[int] = None,
                        audio_aug_prob: dict = None,
                        freq_aug_prob: dict = None,
                        verbose: bool = False,
                        is_bad_audio: bool = True):

        # if length of the lists provided are not equal then raise error.
        if len(path_datasets) != len(names):
            raise ValueError(f'Number of "path_dataset" provided must equal'
                             f'to the number of their respective "names"!')
        if len(path_datasets) != len(n):
            raise ValueError(f'Number of "path_dataset" provided must equal'
                             f'to the number of their respective "n"!')
        if len(names) != len(n):
            raise ValueError(f'Number of "names" provided must equal to the "n"!')

        # if any num in n is less than 1 then raise error
        if any(num < 1 for num in n):
            raise ValueError("list: n must contain 1 or greater value!")

        # if any name in names is not 'real' or 'fake' then raise error.
        if any(name != "real" and name != 'fake' for name in names):
            raise ValueError("list: 'names' must be either 'real' or 'fake' according to the list: 'path_datasets'!")

        tfrecord_writer_real = None
        tfrecord_writer_fake = None

        for path_dataset, name, number in zip(path_datasets, names, n):
            if name == "real":
                # Creating TFRecord first time
                if tfrecord_writer_real is None:
                    real_path = os.path.join(os.getcwd(), 'tfrecords', f'{tfrecord_file_name}_{name}.tfrecords')
                    tfrecord_writer_real = tf_io.TFRecordWriter(real_path)
                self._tfrecord_creator(path_dataset, tfrecord_writer_real, name, number,
                                       verbose=verbose, is_bad_audio=is_bad_audio,
                                       audio_aug_prob=audio_aug_prob, freq_aug_prob=freq_aug_prob)
            elif name == "fake":
                # Creating TFRecord first time
                if tfrecord_writer_fake is None:
                    fake_path = os.path.join(os.getcwd(), 'tfrecords', f'{tfrecord_file_name}_{name}.tfrecords')
                    tfrecord_writer_fake = tf_io.TFRecordWriter(fake_path)
                self._tfrecord_creator(path_dataset, tfrecord_writer_fake, name, number,
                                       verbose=verbose, is_bad_audio=is_bad_audio,
                                       audio_aug_prob=audio_aug_prob, freq_aug_prob=freq_aug_prob)

        tfrecord_writer_fake.close()
        tfrecord_writer_real.close()

    def _tfrecord_creator(self, path_dataset: np.ndarray,
                          tfrecord_writer: tf_io.TFRecordWriter,
                          name: str,
                          n: int = 1,
                          audio_aug_prob: dict = None,
                          freq_aug_prob: dict = None,
                          verbose: bool = False,
                          is_bad_audio: bool = True, ):

        # if n is less than 1 then the main writing loop will not work
        if n < 1:
            raise IndexError(f"n must be greater than 0, in {name}")

        # x and z for keeping track of total files and augmented files, (for debugging reasons)
        x, y, z = 0, 0, 0
        for i, (path, label) in enumerate(path_dataset):
            if verbose:
                print(f'{Fore.BLUE}Loading... {Fore.GREEN}{path} {Fore.BLUE}- {Fore.GREEN}{label}', Fore.RESET)
            x += 1

            feat = None
            # loop for extracting audio and features form audio
            for j in range(n):
                audio = self.load_audio(path, verbose=verbose)
                if audio is None or type(audio) is None:
                    if is_bad_audio:
                        print(f'{Fore.BLUE}Couldn\'t load Path: '
                              f'{Fore.GREEN}{path} {Fore.BLUE}and label: {Fore.GREEN}{label}', Fore.RESET)
                    if not verbose:
                        print(f'\r{Fore.BLUE}{j}:{Fore.CYAN} Failed to load: {Fore.GREEN}{path}', Fore.RESET, end="")
                    break
                # augment audio if range is more than 1 such that the loop, loops more than 1
                if j > 0:
                    # augment every file loaded, n determines the number of time the file will be augmented starting
                    # from 2, so 2 == the file augmented once, 3 == the file augmented twice, and so on.
                    z += 1
                    try:
                        # inherited function
                        audio = self.time_aug(audio, aug_prob=audio_aug_prob)
                    except Exception as e:
                        raise Exception("While augmenting time of audio", type(e).__name__, e)

                try:
                    feat = self.Extractor.extract(audio)  # extracting audio features
                except Exception as e:
                    raise Exception(f"Error name: {type(e).__name__} - {e} while extracting features and normalizing.")

                self.feature_shape = feat.shape

                if j > 0:
                    # augment every file loaded, n determines the number of time the file will be augmented starting
                    # from 2, so 2 == the file augmented once, 3 == the file augmented twice, and so on.
                    y += 1
                    try:
                        # inherited function
                        feat = self.spec_aug(feat, aug_prob=freq_aug_prob)
                    except Exception as e:
                        raise Exception("While augmenting spec of audio", type(e).__name__, e)

                if y != z:
                    raise Exception("While augmenting audio and spectrogram features. (not equally augmented)")

                # calling example function
                example = self._create_examples(audio_features=feat, label=int(label))
                tfrecord_writer.write(example)

            print(f'\r{Fore.BLUE}{i}{Fore.CYAN}: shape{self.feature_shape},'
                  f' type:{type(feat[0, 0])}', Fore.RESET, end="")

        print(f'\nWrote {x + z} files in {name}, {z} files were augmented')

    @staticmethod
    def _create_examples(audio_features: np.ndarray,
                         label: int,
                         verbose: bool = False):

        if audio_features.dtype == np.float32:
            flatten_array = audio_features.flatten()
            audio_bytes = tf_io.serialize_tensor(flatten_array)
            bytes_string_audio = audio_bytes.numpy()
        else:
            raise TypeError(f"'audio_features' content type must be np.float32")

        audio_feature = tf_train.Feature(bytes_list=tf_train.BytesList(value=[bytes_string_audio]))
        label_feature = tf_train.Feature(int64_list=tf_train.Int64List(value=[label]))

        example = tf_train.Example(
            features=tf_train.Features(feature={
                'stft': audio_feature,
                'label': label_feature,
            })
        )
        if verbose:
            print(f'{Fore.BLUE}Creating TFRecord example for label '
                  f'{Fore.GREEN}{audio_features} {Fore.BLUE}- {Fore.GREEN}{label}',
                  Fore.RESET)

        return example.SerializeToString()