data.py

import torch
from torch.utils.data import Dataset

import torchaudio
import torchaudio.transforms as T
from torchaudio.transforms import Resample

from tqdm.auto import tqdm


def normalize_features(features, normalize_type="per_feature"):
    if normalize_type == "per_feature":
        mean = features.mean(dim=-1, keepdim=True)
        std = features.std(dim=-1, keepdim=True)
        normalized_features = (features - mean) / (std + 1e-6)
    elif normalize_type == "per_utterance":
        mean = features.mean()
        std = features.std()
        normalized_features = (features - mean) / (std + 1e-6)
    else:
        raise ValueError("Unsupported normalization type")
    return normalized_features


def get_melspectrogram(
    waveform,
    sample_rate=16000,
    normalize="per_feature",
    window_size=0.025,
    window_stride=0.01,
    window="hann",
    features=80,
    n_fft=512,
    log=True,
    frame_splicing=1,
    dither=0.00001,
    pad_to=0,
    pad_value=0.0,
):
    if waveform.ndim == 1:
        waveform = waveform.unsqueeze(0)  # Shape becomes (1, samples)

    # Apply dithering
    if dither > 0:
        noise = torch.randn_like(waveform) * dither
        waveform += noise

    # Compute the hop length and window length
    hop_length = int(sample_rate * window_stride)
    win_length = int(sample_rate * window_size)

    # Create the MelSpectrogram transform
    mel_transform = T.MelSpectrogram(
        sample_rate=sample_rate,
        n_fft=n_fft,
        win_length=win_length,
        hop_length=hop_length,
        window_fn=torch.hann_window if window == "hann" else None,
        n_mels=features,
    )

    # Compute the mel spectrogram
    melspectrogram = mel_transform(waveform)

    # Logarithmic scaling
    if log:
        melspectrogram = torch.log(melspectrogram + 1e-6)

    # Pad to specified length
    if pad_to > 0:
        pad_length = max(pad_to - melspectrogram.shape[-1], 0)
        melspectrogram = torch.nn.functional.pad(
            melspectrogram, (0, pad_length), value=pad_value
        )

    # Apply frame splicing
    if frame_splicing > 1:
        melspectrogram = melspectrogram.unfold(-1, frame_splicing, frame_splicing).mean(
            dim=-1
        )

    return normalize_features(melspectrogram, normalize)


def process(dataset, split="train", n_feat=80):
    data = []
    for idx, item in tqdm(enumerate(dataset[split]), total=len(dataset[split])):
        try:
            waveform, sample_rate, _, _, _, _ = item

            if sample_rate != 16000:
                resampler = Resample(orig_freq=sample_rate, new_freq=16000)
                waveform = resampler(waveform)
                sample_rate = 16000

            spectrogram = get_melspectrogram(
                waveform, sample_rate=sample_rate, features=n_feat, normalize="per_feature"
            )

            data.append({
                "spectrogram": spectrogram.squeeze(),
            })

        except Exception as e:
            print("err")
            print(str(e))

        if idx == 10_000:
            break

    return data


class BestRQDataset(Dataset):
    def __init__(self, data):
        self.data = data
        self.len = len(self.data)

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        spectrogram = self.data[idx]["spectrogram"]

        return {
            "spectrogram": spectrogram.float(),
        }

    @staticmethod
    def collate_fn(batch):
        spectrograms = [item["spectrogram"] for item in batch]
        lengths = [s.shape[1] for s in spectrograms]

        max_t = max(lengths)
        padded_t = ((max_t + 3) // 4) * 4  # Round up to the nearest multiple of 4

        dim = spectrograms[0].shape[0]

        padded_spectrograms = torch.zeros(len(spectrograms), dim, padded_t)
        for i, s in enumerate(spectrograms):
            padded_spectrograms[i, :, :s.shape[1]] = s.clone().detach()

        return {
            "spectrogram": padded_spectrograms,
            "lengths": torch.tensor(lengths),
        }