tc7/preprocess.py

import math
import torch
import torch.nn as nn
import torchaudio
from torchaudio.transforms import FrequencyMasking
from tja import parse_tja, PyParsingMode
from .config import N_TYPES, SAMPLE_RATE, N_MELS, HOP_LENGTH, TIME_SUB
from .model import TaikoConformer7


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
    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)
    sliding_nps_labels = torch.zeros(
        T_sub, dtype=torch.float32
    )  # New label for sliding NPS

    # 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

        exact_frame_start = t_start.item() * fps

        # Type 1 and 3 are Don, Type 2 and 4 are Ka
        if typ == 1 or typ == 3 or typ == 2 or typ == 4:
            exact_hit_time_sub = exact_frame_start / TIME_SUB

            current_labels = don_labels if (typ == 1 or typ == 3) else ka_labels

            start_points_info = []
            rounded_hit_time_sub = round(exact_hit_time_sub)

            if (
                abs(exact_hit_time_sub - rounded_hit_time_sub) < 1e-6
            ):  # Tolerance for float precision
                idx_single = int(rounded_hit_time_sub)
                if 0 <= idx_single < T_sub:
                    start_points_info.append({"idx": idx_single, "weight": 1.0})
            else:
                idx_floor = math.floor(exact_hit_time_sub)
                idx_ceil = idx_floor + 1

                frac = exact_hit_time_sub - idx_floor
                weight_ceil = frac
                weight_floor = 1.0 - frac

                if weight_floor > 1e-6 and 0 <= idx_floor < T_sub:
                    start_points_info.append({"idx": idx_floor, "weight": weight_floor})
                if weight_ceil > 1e-6 and 0 <= idx_ceil < T_sub:
                    start_points_info.append({"idx": idx_ceil, "weight": weight_ceil})

            for point_info in start_points_info:
                start_idx = point_info["idx"]
                weight = point_info["weight"]
                for k_idx, kernel_val in enumerate(tail_kernel):
                    target_idx = start_idx + k_idx
                    if 0 <= target_idx < T_sub:
                        current_labels[target_idx] = max(
                            current_labels[target_idx].item(),
                            weight * kernel_val.item(),
                        )

        # Type 5, 6, 7 are Drumroll
        elif typ >= 5 and typ <= 7:
            exact_frame_end = t_end.item() * fps
            exact_start_time_sub = exact_frame_start / TIME_SUB
            exact_end_time_sub = exact_frame_end / TIME_SUB

            # Improved drumroll body
            body_loop_start_idx = math.floor(exact_start_time_sub)
            body_loop_end_idx = math.ceil(exact_end_time_sub)

            for dr_idx in range(body_loop_start_idx, body_loop_end_idx):
                if 0 <= dr_idx < T_sub:
                    drumroll_labels[dr_idx] = 1.0

            # Improved drumroll tail (starts from exact_end_time_sub)
            tail_start_points_info = []
            rounded_end_time_sub = round(exact_end_time_sub)
            if abs(exact_end_time_sub - rounded_end_time_sub) < 1e-6:
                idx_single_tail = int(rounded_end_time_sub)
                if 0 <= idx_single_tail < T_sub:
                    tail_start_points_info.append(
                        {"idx": idx_single_tail, "weight": 1.0}
                    )
            else:
                idx_floor_tail = math.floor(exact_end_time_sub)
                idx_ceil_tail = idx_floor_tail + 1

                frac_tail = exact_end_time_sub - idx_floor_tail
                weight_ceil_tail = frac_tail
                weight_floor_tail = 1.0 - frac_tail

                if weight_floor_tail > 1e-6 and 0 <= idx_floor_tail < T_sub:
                    tail_start_points_info.append(
                        {"idx": idx_floor_tail, "weight": weight_floor_tail}
                    )
                if weight_ceil_tail > 1e-6 and 0 <= idx_ceil_tail < T_sub:
                    tail_start_points_info.append(
                        {"idx": idx_ceil_tail, "weight": weight_ceil_tail}
                    )

            for point_info in tail_start_points_info:
                start_idx = point_info["idx"]
                weight = point_info["weight"]
                for k_idx, kernel_val in enumerate(tail_kernel):
                    target_idx = start_idx + k_idx
                    if 0 <= target_idx < T_sub:
                        drumroll_labels[target_idx] = max(
                            drumroll_labels[target_idx].item(),
                            weight * kernel_val.item(),
                        )

    # Calculate sliding window NPS
    note_events = (
        []
    )  # Store tuples of (time_sec, type_is_drumroll_start_or_end, duration_if_drumroll)
    for onset in example[difficulty]:
        typ, t_start_tensor, t_end_tensor, *_ = onset
        t_start = t_start_tensor.item()
        t_end = t_end_tensor.item()

        if typ in [1, 2, 3, 4]:  # Don or Ka
            note_events.append(
                (t_start, False, 0)
            )  # False indicates not a drumroll event, duration 0
        elif typ >= 5 and typ <= 7:  # Drumroll
            note_events.append(
                (t_start, True, t_end - t_start)
            )  # True indicates drumroll start, store duration
            # We don't explicitly need a drumroll end event for this calculation method

    note_events.sort(key=lambda x: x[0])  # Sort by time

    window_duration_seconds = 0.5
    # drumroll_nps_rate = 10.0 # Removed: Will use adaptive rate

    # Step 1: Calculate base_sliding_nps_labels (Don/Ka only)
    base_don_ka_sliding_nps = torch.zeros(T_sub, dtype=torch.float32)
    time_step_duration_sec = TIME_SUB / fps  # Duration of one T_sub segment

    for k_idx in range(T_sub):
        k_window_end_sec = ((k_idx + 1) * TIME_SUB) / fps
        k_window_start_sec = k_window_end_sec - window_duration_seconds

        current_don_ka_count = 0.0
        for event_t, is_drumroll, _ in note_events:
            if not is_drumroll:  # Don or Ka hit
                if k_window_start_sec <= event_t < k_window_end_sec:
                    current_don_ka_count += 1
        base_don_ka_sliding_nps[k_idx] = current_don_ka_count / window_duration_seconds

    # Step 2: Calculate adaptive_drumroll_rates_for_all_events
    adaptive_drumroll_rates_for_all_events = []
    for event_t, is_drumroll, drumroll_dur in note_events:
        if is_drumroll:
            drumroll_start_sec = event_t
            drumroll_end_sec = event_t + drumroll_dur

            slice_start_idx = math.floor(drumroll_start_sec / time_step_duration_sec)
            slice_end_idx = math.ceil(drumroll_end_sec / time_step_duration_sec)

            slice_start_idx = max(0, slice_start_idx)
            slice_end_idx = min(T_sub, slice_end_idx)

            max_nps_in_drumroll_period = 0.0
            if slice_start_idx < slice_end_idx:
                relevant_base_nps_values = base_don_ka_sliding_nps[
                    slice_start_idx:slice_end_idx
                ]
                if relevant_base_nps_values.numel() > 0:
                    max_nps_in_drumroll_period = torch.max(
                        relevant_base_nps_values
                    ).item()

            rate = max(5.0, max_nps_in_drumroll_period)
            adaptive_drumroll_rates_for_all_events.append(rate)
        else:
            adaptive_drumroll_rates_for_all_events.append(0.0)  # Placeholder

    # Step 3: Calculate final sliding_nps_labels using adaptive rates
    # sliding_nps_labels is already initialized with zeros earlier in the function.
    for k_idx in range(T_sub):
        k_window_end_sec = ((k_idx + 1) * TIME_SUB) / fps
        k_window_start_sec = k_window_end_sec - window_duration_seconds

        current_window_total_nps_contribution = 0.0
        for event_idx, (event_t, is_drumroll, drumroll_dur) in enumerate(note_events):
            if is_drumroll:
                drumroll_start_sec = event_t
                drumroll_end_sec = event_t + drumroll_dur

                overlap_start = max(k_window_start_sec, drumroll_start_sec)
                overlap_end = min(k_window_end_sec, drumroll_end_sec)

                if overlap_end > overlap_start:
                    overlap_duration = overlap_end - overlap_start
                    current_adaptive_rate = adaptive_drumroll_rates_for_all_events[
                        event_idx
                    ]
                    current_window_total_nps_contribution += (
                        overlap_duration * current_adaptive_rate
                    )
            else:  # Don or Ka hit
                if k_window_start_sec <= event_t < k_window_end_sec:
                    current_window_total_nps_contribution += (
                        1  # Each hit contributes 1 to the count
                    )

        sliding_nps_labels[k_idx] = (
            current_window_total_nps_contribution / window_duration_seconds
        )

    # Normalize sliding_nps_labels to 0-1 range
    if T_sub > 0:  # Ensure there are elements to normalize
        min_nps_val = torch.min(sliding_nps_labels)
        max_nps_val = torch.max(sliding_nps_labels)
        denominator = max_nps_val - min_nps_val
        if denominator > 1e-6:  # Use a small epsilon for float comparison
            sliding_nps_labels = (sliding_nps_labels - min_nps_val) / denominator
        else:
            # If all values are (nearly) the same
            if max_nps_val > 1e-6:  # If the constant value is positive
                sliding_nps_labels = torch.ones_like(sliding_nps_labels)
            else:  # If the constant value is zero (or very close to it)
                sliding_nps_labels = torch.zeros_like(sliding_nps_labels)

    duration_seconds = wav_tensor.shape[-1] / SAMPLE_RATE
    nps = num_valid_notes / duration_seconds if duration_seconds > 0 else 0.0

    parsed = parse_tja(example["tja"], mode=PyParsingMode.Full)
    chart = next(
        (chart for chart in parsed.charts if chart.course.lower() == difficulty), None
    )
    difficulty_id = (
        0
        if difficulty == "easy"
        else (
            1
            if difficulty == "normal"
            else 2 if difficulty == "hard" else 3 if difficulty == "oni" else 4
        )  # Assuming 4 for edit/ura
    )
    level = chart.level if chart else 0

    # --- CNN shape inference and label padding/truncation ---
    # Simulate CNN to get output time length (T_cnn)
    dummy_model = TaikoConformer7()
    with torch.no_grad():
        cnn_out = dummy_model.cnn(mel.unsqueeze(0))  # (1, C, F, T_cnn)
        _, _, _, T_cnn = cnn_out.shape

    # Pad or truncate labels to T_cnn
    def pad_or_truncate(label, out_len):
        if label.shape[0] < out_len:
            pad = torch.zeros(out_len - label.shape[0], dtype=label.dtype)
            return torch.cat([label, pad], dim=0)
        else:
            return label[:out_len]

    don_labels = pad_or_truncate(don_labels, T_cnn)
    ka_labels = pad_or_truncate(ka_labels, T_cnn)
    drumroll_labels = pad_or_truncate(drumroll_labels, T_cnn)
    sliding_nps_labels = pad_or_truncate(sliding_nps_labels, T_cnn)  # Pad new label

    # For conformer input lengths: this should be T_cnn
    conformer_sequence_length = T_cnn  # This is the actual sequence length after CNN

    print(
        f"Processed {num_valid_notes} notes in {duration_seconds:.2f} seconds, NPS: {nps:.2f}, Difficulty: {difficulty_id}, Level: {level}"
    )

    return {
        "mel": mel,  # (1, N_MELS, T)
        "don_labels": don_labels,  # (T_cnn,)
        "ka_labels": ka_labels,  # (T_cnn,)
        "drumroll_labels": drumroll_labels,  # (T_cnn,)
        "sliding_nps_labels": sliding_nps_labels,  # Add new label (T_cnn,)
        "nps": torch.tensor(nps, dtype=torch.float32),
        "difficulty": torch.tensor(difficulty_id, dtype=torch.long),
        "level": torch.tensor(level, dtype=torch.long),
        "duration_seconds": torch.tensor(duration_seconds, dtype=torch.float32),
        "length": torch.tensor(
            conformer_sequence_length, dtype=torch.long
        ),  # Use T_cnn for conformer and loss masking
    }


def collate_fn(batch):
    mels_list = [b["mel"].squeeze(0).transpose(0, 1) for b in batch]  # (T, N_MELS)
    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]
    sliding_nps_labels_list = [b["sliding_nps_labels"] for b in batch]  # New label list
    nps_list = [b["nps"] for b in batch]
    difficulty_list = [b["difficulty"] for b in batch]
    level_list = [b["level"] for b in batch]
    durations_list = [b["duration_seconds"] for b in batch]
    lengths_list = [b["length"] for b in batch]  # These are T_cnn_i for each example

    # Pad mels
    padded_mels = nn.utils.rnn.pad_sequence(
        mels_list, batch_first=True
    )  # (B, T_max_mel, N_MELS)
    reshaped_mels = padded_mels.transpose(1, 2).unsqueeze(1)
    # T_max_mel_batch = padded_mels.shape[1] # Max mel length in batch, not used for label padding anymore

    # Determine max sequence length for labels (max T_cnn in batch)
    max_label_len = 0
    if lengths_list:  # handle empty batch case
        max_label_len = max(l.item() for l in lengths_list) if lengths_list else 0

    # Pad labels to max_label_len (max_t_cnn_in_batch)
    def pad_label_to_max_len(label_tensor, target_len):
        current_len = label_tensor.shape[0]
        if current_len < target_len:
            padding_size = target_len - current_len
            # Ensure padding is created on the same device as the label_tensor
            padding = torch.zeros(
                padding_size, dtype=label_tensor.dtype, device=label_tensor.device
            )
            return torch.cat((label_tensor, padding), dim=0)
        elif (
            current_len > target_len
        ):  # Should ideally not happen if lengths_list is correct
            return label_tensor[:target_len]
        return label_tensor

    don_labels = torch.stack(
        [pad_label_to_max_len(l, max_label_len) for l in don_labels_list]
    )
    ka_labels = torch.stack(
        [pad_label_to_max_len(l, max_label_len) for l in ka_labels_list]
    )
    drumroll_labels = torch.stack(
        [pad_label_to_max_len(l, max_label_len) for l in drumroll_labels_list]
    )
    sliding_nps_labels = torch.stack(
        [pad_label_to_max_len(l, max_label_len) for l in sliding_nps_labels_list]
    )  # Pad new labels

    actual_lengths = torch.tensor([l.item() for l in lengths_list], dtype=torch.long)

    return {
        "mel": reshaped_mels,
        "don_labels": don_labels,
        "ka_labels": ka_labels,
        "drumroll_labels": drumroll_labels,
        "sliding_nps_labels": sliding_nps_labels,  # Add new batched labels
        "lengths": actual_lengths,  # for conformer and loss masking (T_cnn_i for each item)
        "nps": torch.stack(nps_list),
        "difficulty": torch.stack(difficulty_list),
        "level": torch.stack(level_list),
        "durations": torch.stack(durations_list),
    }