amt/src/utils/preprocess/preprocess_urmp.py

"""preprocess_mir1k.py"""
import os
import shutil
import glob
import re
import json
from typing import Dict, List, Tuple
import numpy as np
from utils.audio import get_audio_file_info, load_audio_file
from utils.midi import midi2note, note_event2midi
from utils.note2event import note2note_event, mix_notes, sort_notes, validate_notes, trim_overlapping_notes
from utils.event2note import event2note_event
from utils.note_event_dataclasses import Note, NoteEvent
from utils.utils import note_event2token2note_event_sanity_check, freq_to_midi

MT3_TEST_IDS = [1, 2, 12, 13, 24, 25, 31, 38, 39]
PROGRAM_STR2NUM = {
    'vn': 40,
    'va': 41,
    'vc': 42,
    'db': 43,
    'fl': 73,
    'ob': 68,
    'cl': 71,
    'sax': 65,  # The type of sax used in the dataset is not clear. We guess it would be alto sax.
    'bn': 70,
    'tpt': 56,
    'hn': 60,  # Just annotated as horn. We guess it would be french horn, due to the pitch range.
    'tbn': 57,
    'tba': 58,
}


def delete_hidden_files(base_dir):
    for hidden_file in glob.glob(os.path.join(base_dir, '**/.*'), recursive=True):
        os.remove(hidden_file)
        print(f"Deleted: {hidden_file}")


def convert_annotation_to_notes(id, program, ann_files):
    notes = []
    for ann_file, prog in zip(ann_files, program):
        data = np.loadtxt(ann_file)
        onset = data[:, 0]
        freq = data[:, 1]
        duration = data[:, 2]

        notes_by_instr = []
        for o, f, d in zip(onset, freq, duration):
            notes_by_instr.append(
                Note(
                    is_drum=False,
                    program=prog,
                    onset=o,
                    offset=o + d,
                    pitch=freq_to_midi(f),
                    velocity=1))
        notes = mix_notes([notes, notes_by_instr], sort=True, trim_overlap=True, fix_offset=True)
    notes = sort_notes(notes)
    note_events = note2note_event(notes, sort=True)
    duration_sec = note_events[-1].time + 0.01
    return {  # notes
            'urmp_id': id,
            'program': program,
            'is_drum': [0] * len(program),
            'duration_sec': duration_sec,
            'notes': notes,
        }, {  # note_events
            'guitarset_id': id,
            'program': program,
            'is_drum': [0] * len(program),
            'duration_sec': duration_sec,
            'note_events': note_events,
        }


def create_audio_stem(audio_tracks, id, program, n_frames):
    max_length = max([len(tr) for tr in audio_tracks])
    max_length = max(max_length, n_frames)
    n_tracks = len(audio_tracks)
    audio_array = np.zeros((n_tracks, max_length), dtype=np.float16)
    for j, audio in enumerate(audio_tracks):
        audio_array[j, :len(audio)] = audio

    return {
        'urmp_id': id,
        'program': np.array(program),
        'is_drum': np.array([0] * len(program), dtype=np.int64),
        'n_frames': n_frames,  # int
        'audio_array': audio_array  # (n_tracks, n_frames)
    }


def data_bug_fix(base_dir):
    files = glob.glob(os.path.join(base_dir, '15_Surprise_tpt_tpt_tbn', '*3_tpt*.*'))
    for file in files:
        new_file = file.replace('3_tpt', '3_tbn')
        shutil.move(file, new_file)
        print(f"Renamed: {file} -> {new_file}")


def preprocess_urmp16k(data_home=os.PathLike,
                       dataset_name='urmp',
                       delete_source_files: bool = False,
                       sanity_check=True) -> None:
    """
    URMP dataset does not have official split information. We follow the split used in MT3 paper.

    About:
    - 44 pieces of classical music
    - Duet, Trio, Quartet, Quintet of strings or winds or mixed
    - Multi-stem audio 
    - MIDI file is unaligned, it is for score
    - Annotation (10ms hop) is provided.
    - There is timing issue for annotation
    - We do not use video

    Splits:
        - train: 35 files, following MT3
        - test: 9 files, follwing MT3
        - all: 44 files

    Writes:
        - {dataset_name}_{split}_file_list.json: a dictionary with the following keys:
        {
            index:
            {
                'urmp_id': urmp_id,
                'n_frames': (int),
                'stem_file': 'path/to/stem.npy',
                'mix_audio_file': 'path/to/mix.wav',
                'notes_file': 'path/to/notes.npy',
                'note_events_file': 'path/to/note_events.npy',
                'midi_file': 'path/to/midi.mid', # this is 120bpm converted midi file from note_events
                'program': List[int], #  
                'is_drum': List[int], # [0] or [1]
            }
        }
    """

    # Directory and file paths
    base_dir = os.path.join(data_home, dataset_name + '_yourmt3_16k')
    output_index_dir = os.path.join(data_home, 'yourmt3_indexes')
    os.makedirs(output_index_dir, exist_ok=True)

    # Databug fix
    data_bug_fix(base_dir)

    # Delete hidden files
    delete_hidden_files(base_dir)

    # Create file list for split==all
    file_list = dict()
    for dir_name in sorted(os.listdir(base_dir)):
        if dir_name.startswith('.'):
            continue
        if 'Supplementary' in dir_name:
            continue
        # urmp_id
        id = dir_name.split('_')[0]
        title = dir_name.split('_')[1]

        # program
        program_strings = dir_name.split('_')[2:]
        program = [PROGRAM_STR2NUM[p] for p in program_strings]

        # is_drum
        is_drum = [0] * len(program)

        # file paths
        stem_file = os.path.join(base_dir, dir_name, 'stem.npy')
        mix_audio_file = glob.glob(os.path.join(base_dir, dir_name, 'AuMix*.wav'))[0]
        notes_file = os.path.join(base_dir, dir_name, 'notes.npy')
        note_events_file = os.path.join(base_dir, dir_name, 'note_events.npy')
        midi_file = os.path.join(base_dir, dir_name, f'{str(id)}_120bpm_converted.mid')

        # n_frames
        fs, n_frames, n_channels = get_audio_file_info(mix_audio_file)
        assert fs == 16000 and n_channels == 1

        # Fill out a file list
        file_list[id] = {
            'urmp_id': id,
            'n_frames': n_frames,
            'stem_file': stem_file,
            'mix_audio_file': mix_audio_file,
            'notes_file': notes_file,
            'note_events_file': note_events_file,
            'midi_file': midi_file,
            'program': program,
            'is_drum': is_drum,
        }

        # Process Annotations
        ann_files = [
            os.path.join(base_dir, dir_name, f'Notes_{i+1}_{p}_{str(id)}_{title}.txt')
            for i, p in enumerate(program_strings)
        ]

        # Check if all files exist
        for ann_file in ann_files:
            assert os.path.exists(ann_file), f"{ann_file} does not exist."
        assert len(program) == len(ann_files)

        # Create and save notes and note_events from annotation
        notes, note_events = convert_annotation_to_notes(id, program, ann_files)
        np.save(notes_file, notes, allow_pickle=True, fix_imports=False)
        print(f'Created {notes_file}')
        np.save(note_events_file, note_events, allow_pickle=True, fix_imports=False)
        print(f'Created {note_events_file}')

        # Create 120bpm MIDI file from note_events
        note_event2midi(note_events['note_events'], midi_file)
        print(f'Created {midi_file}')

        # Process Audio
        audio_tracks = []
        for i, p in enumerate(program_strings):
            audio_sep_file = os.path.join(base_dir, dir_name, f'AuSep_{i+1}_{p}_{id}_{title}.wav')
            audio_track = load_audio_file(audio_sep_file, dtype=np.int16) / 2**15  # returns bytes
            audio_tracks.append(audio_track.astype(np.float16))
            if delete_source_files:
                os.remove(audio_sep_file)

        stem_content = create_audio_stem(audio_tracks, id, program, n_frames)
        np.save(stem_file, stem_content, allow_pickle=True, fix_imports=False)
        print(f'Created {stem_file}')

        # Sanity check
        if sanity_check:
            recon_notes, _ = midi2note(midi_file)
            recon_note_events = note2note_event(recon_notes)
            note_event2token2note_event_sanity_check(recon_note_events, notes['notes'])

        # File existence check
        assert os.path.exists(mix_audio_file)

    # Create index for splits
    file_list_all = {}
    for i, key in enumerate(file_list.keys()):
        file_list_all[i] = file_list[key]

    file_list_train = {}
    i = 0
    for key in file_list.keys():
        if int(key) not in MT3_TEST_IDS:
            file_list_train[i] = file_list[key]
            i += 1

    file_list_test = {}
    i = 0
    for key in file_list.keys():
        if int(key) in MT3_TEST_IDS:
            file_list_test[i] = file_list[key]
            i += 1

    all_fl = {'all': file_list_all, 'train': file_list_train, 'test': file_list_test}

    # Save index
    for split in ['all', 'train', 'test']:
        output_index_file = os.path.join(output_index_dir, f'{dataset_name}_{split}_file_list.json')
        with open(output_index_file, 'w') as f:
            json.dump(all_fl[split], f, indent=4)
        print(f'Created {output_index_file}')