tc5/preprocess.py

import math
import torch
import torch.nn as nn
import torchaudio
from torchaudio.transforms import FrequencyMasking
from .config import N_TYPES, SAMPLE_RATE, N_MELS, HOP_LENGTH, TIME_SUB
from .model import TaikoConformer5


mel_transform = torchaudio.transforms.MelSpectrogram(
    sample_rate=SAMPLE_RATE,
    n_mels=N_MELS,
    hop_length=HOP_LENGTH,
    n_fft=2048,
)


freq_mask = FrequencyMasking(freq_mask_param=15)


def preprocess(example, difficulty="oni"):
    wav_tensor = example["audio"]["array"]
    sr = example["audio"]["sampling_rate"]
    # 1) load & resample
    if sr != SAMPLE_RATE:
        wav_tensor = torchaudio.functional.resample(wav_tensor, sr, SAMPLE_RATE)
    # normalize audio
    wav_tensor = wav_tensor / (wav_tensor.abs().max() + 1e-8)
    # add random Gaussian noise
    if torch.rand(1).item() < 0.5:
        wav_tensor = wav_tensor + 0.005 * torch.randn_like(wav_tensor)
    # 2) mel: (1, N_MELS, T)
    mel = mel_transform(wav_tensor).unsqueeze(0)
    # apply SpecAugment
    # we don't use time masking since we don't want model to predict notes when they are masked
    mel = freq_mask(mel)
    _, _, T = mel.shape
    # 3) build label sequence of length ceil(T / TIME_SUB)
    T_sub = math.ceil(T / TIME_SUB)

    # Initialize energy-based labels for Don, Ka, Drumroll
    don_labels = torch.zeros(T_sub, dtype=torch.float32)
    ka_labels = torch.zeros(T_sub, dtype=torch.float32)
    drumroll_labels = torch.zeros(T_sub, dtype=torch.float32)

    # Define exponential decay tail parameters
    tail_length = 40  # number of frames for decay tail
    decay_rate = 8.0  # decay rate parameter, adjust as needed
    tail_kernel = torch.exp(
        -torch.arange(0, tail_length, dtype=torch.float32) / decay_rate
    )

    fps = SAMPLE_RATE / HOP_LENGTH
    num_valid_notes = 0
    for onset in example[difficulty]:
        typ, t_start, t_end, *_ = onset

        # Assuming N_TYPES in config is appropriately set (e.g., 7 or more)
        if typ < 1 or typ > N_TYPES:  # Filter out invalid types
            continue

        num_valid_notes += 1

        f = int(round(t_start.item() * fps))
        idx = f // TIME_SUB
        if 0 <= idx < T_sub:
            # Apply exponential decay kernel to the corresponding energy channel
            # Type 1 and 3 are Don
            if typ == 1 or typ == 3:
                for i, val in enumerate(tail_kernel):
                    target_idx = idx + i
                    if 0 <= target_idx < T_sub:
                        don_labels[target_idx] = max(
                            don_labels[target_idx].item(), val.item()
                        )
            # Type 2 and 4 are Ka
            elif typ == 2 or typ == 4:
                for i, val in enumerate(tail_kernel):
                    target_idx = idx + i
                    if 0 <= target_idx < T_sub:
                        ka_labels[target_idx] = max(
                            ka_labels[target_idx].item(), val.item()
                        )
            # Type 5, 6, 7 are Drumroll
            elif typ >= 5 and typ <= 7:
                f_end = int(round(t_end.item() * fps))
                idx_end = f_end // TIME_SUB

                for dr in range(idx, idx_end):
                    if 0 <= dr < T_sub:
                        drumroll_labels[dr] = 1.0

                for i, val in enumerate(tail_kernel):
                    target_idx = idx_end + i
                    if 0 <= target_idx < T_sub:
                        drumroll_labels[target_idx] = max(
                            drumroll_labels[target_idx].item(), val.item()
                        )

    duration_seconds = wav_tensor.shape[-1] / SAMPLE_RATE
    nps = num_valid_notes / duration_seconds if duration_seconds > 0 else 0.0
    print(
        f"Processed {num_valid_notes} notes in {duration_seconds:.2f} seconds, NPS: {nps:.2f}"
    )

    return {
        "mel": mel,
        "don_labels": don_labels,
        "ka_labels": ka_labels,
        "drumroll_labels": drumroll_labels,
        "nps": torch.tensor(nps, dtype=torch.float32),
        "duration_seconds": torch.tensor(duration_seconds, dtype=torch.float32),
    }


def collate_fn(batch):
    mels_list = [b["mel"].squeeze(0).transpose(0, 1) for b in batch]  # (T, N_MELS)
    # Extract new energy-based labels
    don_labels_list = [b["don_labels"] for b in batch]
    ka_labels_list = [b["ka_labels"] for b in batch]
    drumroll_labels_list = [b["drumroll_labels"] for b in batch]

    nps_list = [b["nps"] for b in batch]  # Extract NPS
    durations_list = [b["duration_seconds"] for b in batch]  # Extract durations

    # Pad mels
    padded_mels = nn.utils.rnn.pad_sequence(
        mels_list, batch_first=True
    )  # (B, T_max, N_MELS)
    # Reshape for CNN: (B, 1, N_MELS, T_max)
    reshaped_mels = padded_mels.transpose(1, 2).unsqueeze(1)

    # Simulate CNN time downsampling to get output lengths
    dummy_model_for_shape_inference = TaikoConformer5()
    dummy_cnn = dummy_model_for_shape_inference.cnn
    with torch.no_grad():
        cnn_out = dummy_cnn(reshaped_mels)  # Use reshaped_mels that has batch dim
        _, _, _, T_cnn = cnn_out.shape

    padded_don_labels = []
    padded_ka_labels = []
    padded_drumroll_labels = []
    # lengths = [] # This was for original presence/type labels, conformer_input_lengths is used for model

    for i in range(len(batch)):
        d_labels = don_labels_list[i]
        k_labels = ka_labels_list[i]
        dr_labels = drumroll_labels_list[i]

        item_original_T_sub = d_labels.shape[
            0
        ]  # Assuming all label types have same original length
        out_len = T_cnn  # Target length for labels is T_cnn

        # Pad or truncate don_labels
        if item_original_T_sub < out_len:
            pad_d = torch.full(
                (out_len - item_original_T_sub,),
                0,  # Pad with 0 for energy labels
                dtype=d_labels.dtype,
                device=d_labels.device,
            )
            padded_d = torch.cat([d_labels, pad_d], dim=0)
        else:
            padded_d = d_labels[:out_len]
        padded_don_labels.append(padded_d)

        # Pad or truncate ka_labels
        if item_original_T_sub < out_len:
            pad_k = torch.full(
                (out_len - item_original_T_sub,),
                0,  # Pad with 0 for energy labels
                dtype=k_labels.dtype,
                device=k_labels.device,
            )
            padded_k = torch.cat([k_labels, pad_k], dim=0)
        else:
            padded_k = k_labels[:out_len]
        padded_ka_labels.append(padded_k)

        # Pad or truncate drumroll_labels
        if item_original_T_sub < out_len:
            pad_dr = torch.full(
                (out_len - item_original_T_sub,),
                0,  # Pad with 0 for energy labels
                dtype=dr_labels.dtype,
                device=dr_labels.device,
            )
            padded_dr = torch.cat([dr_labels, pad_dr], dim=0)
        else:
            padded_dr = dr_labels[:out_len]
        padded_drumroll_labels.append(padded_dr)

    # For Conformer input lengths: lengths of mel sequences after CNN subsampling
    # (Assuming CNN does not subsample in time, T_cnn is effectively T_mel_padded)
    # The `lengths` for the Conformer should be based on the mel input to the conformer part.
    # The existing calculation for conformer_input_lengths seems to relate to TIME_SUB.
    # If the Conformer input itself is not subsampled by TIME_SUB, this might need review.
    # For now, keeping the existing conformer_input_lengths logic as it's outside the scope of label change.
    conformer_input_lengths = [
        math.ceil(mels_list[i].shape[0] / TIME_SUB) for i in range(len(batch))
    ]
    conformer_input_lengths = torch.tensor(
        [min(l, T_cnn) for l in conformer_input_lengths], dtype=torch.long
    )

    return {
        "mel": reshaped_mels,
        "don_labels": torch.stack(padded_don_labels),
        "ka_labels": torch.stack(padded_ka_labels),
        "drumroll_labels": torch.stack(padded_drumroll_labels),
        "lengths": conformer_input_lengths,  # These are for the Conformer model
        "nps": torch.stack(nps_list),
        "durations": torch.stack(durations_list),
    }