amt/src/utils/datasets_eval.py

# Copyright 2024 The YourMT3 Authors.
#
# 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
#
# Please see the details in the LICENSE file.
import json
import os
from typing import Dict, Any, Union, Tuple, Optional

import torch
import numpy as np
from einops import rearrange
from torch.utils.data import DataLoader, Dataset
from utils.audio import load_audio_file, slice_padded_array
from utils.tokenizer import EventTokenizerBase, NoteEventTokenizer
from utils.note2event import slice_multiple_note_events_and_ties_to_bundle
from utils.note_event_dataclasses import Note, NoteEvent, NoteEventListsBundle
from utils.task_manager import TaskManager
from config.config import shared_cfg
from config.config import audio_cfg as default_audio_cfg

UNANNOTATED_PROGRAM = 129


class AudioFileDataset(Dataset):
    """
    🎧 AudioFileDataset for validation/test:
    
    This dataset class is designed to be used ONLY with `batch_size=None` and 
    returns sliced audio segments and unsliced notes and sliced note events for
     a single song when `__getitem__` is called.

    Args:
        file_list (Union[str, bytes, os.PathLike], optional):
            Path to the file list. e.g. "../../data/yourmt3_indexes/slakh_validation_file_list.json"
        task_manager (TaskManager, optional): TaskManager instance. Defaults to TaskManager().
        fs (int, optional): Sampling rate. Defaults to 16000.
        seg_len_frame (int, optional): Segment length in frames. Defaults to 32767.
        seg_hop_frame (int, optional): Segment hop in frames. Defaults to 32767.
        sub_batch_size (int, optional): Sub-batch size that will be used in 
            generation of tokens. Defaults to 32.
        max_num_files (int, optional): Maximum number of files to be loaded. Defaults to None.
        
    
    Variables:
        file_list:
            '{dataset_name}_{split}_file_list.json' has the following keys:
            {
                'index':
                    {
                        'mtrack_id': mtrack_id,
                        'n_frames': n of audio frames
                        'stem_file': Dict of stem audio file info
                        'mix_audio_file': mtrack.mix_path,
                        'notes_file': available only for 'validation' and 'test'
                        'note_events_file': available only for 'train' and 'validation'
                        'midi_file': mtrack.midi_path
                    }
            }
            
    __getitem__(index) returns:

        audio_segment:
            torch.FloatTensor: (nearest_N_divisable_by_sub_batch_size, 1, seg_len_frame)

        notes_dict:
            {
                'mtrack_id': int,
                'program': List[int],
                'is_drum': bool, 
                'duration_sec': float, 
                'notes': List[Note], 
            }
            
        token_array:
            torch.LongTensor: (n_segments, seg_len_frame)

    """

    def __init__(
            self,
            file_list: Union[str, bytes, os.PathLike],
            task_manager: TaskManager = TaskManager(),
            #  tokenizer: Optional[EventTokenizerBase] = None,
            fs: int = 16000,
            seg_len_frame: int = 32767,
            seg_hop_frame: int = 32767,
            max_num_files: Optional[int] = None) -> None:

        # load the file list
        with open(file_list, 'r') as f:
            fl = json.load(f)
        file_list = {int(key): value for key, value in fl.items()}
        if max_num_files:  # reduce the number of files
            self.file_list = dict(list(file_list.items())[:max_num_files])
        else:
            self.file_list = file_list

        self.fs = fs
        self.seg_len_frame = seg_len_frame
        self.seg_len_sec = seg_len_frame / fs
        self.seg_hop_frame = seg_hop_frame
        self.task_manager = task_manager

    def __getitem__(self, index: int) -> Tuple[np.ndarray, Dict, NoteEventListsBundle]:
        # get metadata
        metadata = self.file_list[index]
        audio_file = metadata['mix_audio_file']
        notes_file = metadata['notes_file']
        note_events_file = metadata['note_events_file']

        # load the audio
        audio = load_audio_file(audio_file, dtype=np.int16)  # returns bytes
        audio = audio / 2**15
        audio = audio.astype(np.float32)
        audio = audio.reshape(1, -1)
        audio_segments = slice_padded_array(
            audio,
            self.seg_len_frame,
            self.seg_hop_frame,
            pad=True,
        )  # (n_segs, seg_len_frame)
        audio_segments = rearrange(audio_segments, 'n t -> n 1 t').astype(np.float32)
        num_segs = audio_segments.shape[0]

        # load all notes and from a file (of a single song)
        notes_dict = np.load(notes_file, allow_pickle=True, fix_imports=False).tolist()

        # TODO: add midi_file path in preprocessing instead of here
        notes_dict['midi_file'] = metadata['midi_file']

        # tokenize note_events
        note_events_dict = np.load(note_events_file, allow_pickle=True, fix_imports=False).tolist()

        if self.task_manager.tokenizer is not None:
            # not using seg_len_sec to avoid accumulated rounding errors
            start_times = [i * self.seg_hop_frame / self.fs for i in range(num_segs)]
            note_event_segments = slice_multiple_note_events_and_ties_to_bundle(
                note_events_dict['note_events'],
                start_times,
                self.seg_len_sec,
            )

            # Support for multi-channel decoding
            if UNANNOTATED_PROGRAM in notes_dict['program']:
                has_unannotated_segments = [True] * num_segs
            else:
                has_unannotated_segments = [False] * num_segs

            token_array = self.task_manager.tokenize_note_events_batch(note_event_segments,
                                                                       start_time_to_zero=False,
                                                                       sort=True)
            # note_token_array = self.task_manager.tokenize_note_events_batch(note_event_segments,
            #                                                                 start_time_to_zero=False,
            #                                                                 sort=True)
            # task_token_array = self.task_manager.tokenize_task_events_batch(note_event_segments,
            #                                                                 has_unannotated_segments)

        # Shape:
        #   processed_audio_array: (num_segs, 1, nframe)
        #   notes_dict: Dict
        #   note_token_array: (num_segs, decoding_ch, max_note_token_len)
        #   task_token_array: (num_segs, decoding_ch, max_task_token_len)
        # return torch.from_numpy(audio_segments), notes_dict, torch.from_numpy(
        #     note_token_array).long(), torch.from_numpy(task_token_array).long()
        return torch.from_numpy(audio_segments), notes_dict, torch.from_numpy(token_array).long()

        # # Tokenize/pad note_event_segments -> array of token and mask
        # max_len = self.tokenizer.max_length
        # token_array = np.zeros((num_segs, max_len), dtype=np.int32)

        # for i, tup in enumerate(list(zip(*note_event_segments.values()))):
        #     padded_tokens = self.tokenizer.encode_plus(*tup)
        #     token_array[i, :] = padded_tokens
        # return torch.from_numpy(audio_segments), notes_dict, torch.from_numpy(token_array).long()

    def __len__(self) -> int:
        return len(self.file_list)


def get_eval_dataloader(
    dataset_name: str,
    split: str = 'validation',
    dataloader_config: Dict = {"num_workers": 0},
    task_manager: TaskManager = TaskManager(),
    # tokenizer: Optional[EventTokenizerBase] = NoteEventTokenizer('mt3'),
    max_num_files: Optional[int] = None,
    audio_cfg: Optional[Dict] = None,
) -> DataLoader:
    """
    🎧 get_audio_file_dataloader:
    
    This function returns a dataloader for AudioFileDataset that returns padded slices
    of audio samples with the divisable number of sub-batch size.
    """
    data_home = shared_cfg["PATH"]["data_home"]
    file_list = f"{data_home}/yourmt3_indexes/{dataset_name}_{split}_file_list.json"

    if audio_cfg is None:
        audio_cfg = default_audio_cfg

    ds = AudioFileDataset(
        file_list,
        task_manager=task_manager,
        # tokenizer=tokenizer,
        seg_len_frame=int(audio_cfg["input_frames"]),  # Default: 32767
        seg_hop_frame=int(audio_cfg["input_frames"]),  # Default: 32767
        max_num_files=max_num_files)
    dl = DataLoader(ds, batch_size=None, collate_fn=lambda k: k, **dataloader_config)
    return dl