Spaces:
Sleeping
Sleeping
| # 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 | |