""" preprocess_mtrack_slakh.py
"""
import os
import time
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
from utils.note2event import note2note_event, mix_notes
import mirdata
from utils.mirdata_dev.datasets import slakh16k
def create_audio_stem_from_mtrack(ds: mirdata.core.Dataset,
mtrack_id: str,
delete_source_files: bool = False) -> Dict:
"""Extracts audio stems and metadata from a multitrack."""
mtrack = ds.multitrack(mtrack_id)
track_ids = mtrack.track_ids
max_length = 0
program_numbers = []
is_drum = []
audio_tracks = [] # multi-channel audio array (C, T)
# collect all the audio tracks and their metadata
for track_id in track_ids:
track = ds.track(track_id)
audio_file = track.audio_path
program_numbers.append(track.program_number)
is_drum.append(1) if track.is_drum else is_drum.append(0)
fs, n_frames, n_channels = get_audio_file_info(audio_file)
assert (fs == 16000 and n_channels == 1)
max_length = n_frames if n_frames > max_length else max_length
audio = load_audio_file(audio_file, dtype=np.int16) # returns bytes
audio = audio / 2**15
audio = audio.astype(np.float16)
audio_tracks.append(audio)
if delete_source_files:
print(f'🗑️ Deleting {audio_file} ...')
os.remove(audio_file)
# collate all the audio tracks into a single array
n_tracks = len(track_ids)
audio_array = np.zeros((n_tracks, max_length), dtype=np.float16)
for j, audio in enumerate(audio_tracks):
audio_array[j, :len(audio)] = audio
stem_content = {
'mtrack_id': mtrack_id, # str
'program': np.array(program_numbers, dtype=np.int64),
'is_drum': np.array(is_drum, dtype=np.int64),
'n_frames': max_length, # int
'audio_array': audio_array # (n_tracks, n_frames)
}
return stem_content
def create_note_event_and_note_from_mtrack_mirdata(
ds: mirdata.core.Dataset,
mtrack_id: str,
fix_bass_octave: bool = True) -> Tuple[Dict, Dict]:
"""Extracts note or note_event and metadata from a multitrack:
Args:
ds (mirdata.core.Dataset): Slakh dataset.
mtrack_id (str): multitrack id.
Returns:
notes (dict): note events and metadata.
note_events (dict): note events and metadata.
"""
mtrack = ds.multitrack(mtrack_id)
track_ids = mtrack.track_ids
program_numbers = []
is_drum = []
mixed_notes = []
duration_sec = 0.
# mix notes from all stem midi files
for track_id in track_ids:
track = ds.track(track_id)
stem_midi_file = track.midi_path
notes, dur_sec = midi2note(
stem_midi_file,
binary_velocity=True,
ch_9_as_drum=False, # checked safe to set to False in Slakh
force_all_drum=True if track.is_drum else False,
force_all_program_to=None, # Slakh always has program number
trim_overlap=True,
fix_offset=True,
quantize=True,
verbose=0,
minimum_offset_sec=0.01,
drum_offset_sec=0.01)
if fix_bass_octave == True and track.program_number in np.arange(32, 40):
if track.plugin_name == 'scarbee_jay_bass_slap_both.nkm':
pass
else:
for note in notes:
note.pitch -= 12
print("Fixed bass octave for track", track_id)
mixed_notes = mix_notes((mixed_notes, notes), True, True, True)
program_numbers.append(track.program_number)
is_drum.append(1) if track.is_drum else is_drum.append(0)
duration_sec = max(duration_sec, dur_sec)
# convert mixed notes to note events
mixed_note_events = note2note_event(mixed_notes, sort=True, return_activity=True)
return { # notes
'mtrack_id': mtrack_id, # str
'program': np.array(program_numbers, dtype=np.int64), # (n,)
'is_drum': np.array(is_drum, dtype=np.int64), # (n,) with 1 is drum
'duration_sec': duration_sec, # float
'notes': mixed_notes # list of Note instances
}, { # note_events
'mtrack_id': mtrack_id, # str
'program': np.array(program_numbers, dtype=np.int64), # (n,)
'is_drum': np.array(is_drum, dtype=np.int64), # (n,) with 1 is drum
'duration_sec': duration_sec, # float
'note_events': mixed_note_events # list of NoteEvent instances
}
def preprocess_slakh16k(data_home: str,
run_checksum: bool = False,
delete_source_files: bool = False,
fix_bass_octave: bool = True) -> None:
"""
Processes the Slakh dataset and extracts stems for each multitrack.
Args:
data_home (str): path to the Slakh data.
run_checksum (bool): if True, validates the dataset using its checksum. Default is False.
delete_source_files (bool): if True, deletes original audio files. Default is False.
fix_bass_octave (bool): if True, fixes the bass to be -1 octave. Slakh bass is annotated as +1 octave. Default is True.
Writes:
- {dataset_name}_{split}_file_list.json: a dictionary with the following keys:
{
'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
}
"""
start_time = time.time()
ds = slakh16k.Dataset(data_home=data_home, version='2100-yourmt3-16k')
if run_checksum:
print('Checksum for slakh dataset...')
ds.validate()
print('Preprocessing slakh dataset...')
mtrack_split_dict = ds.get_mtrack_splits()
for split in ['train', 'validation', 'test']:
file_list = {} # write a file list for each split
mtrack_ids = mtrack_split_dict[split]
for i, mtrack_id in enumerate(mtrack_ids):
print(f'🏃🏻♂️: processing {mtrack_id} ({i+1}/{len(mtrack_ids)} in {split})')
mtrack = ds.multitrack(mtrack_id)
output_dir = os.path.dirname(mtrack.mix_path) # same as mtrack
"""Audio: get stems (as array) and metadata from the multitrack"""
stem_content = create_audio_stem_from_mtrack(ds, mtrack_id, delete_source_files)
# save the audio array and metadata to disk
stem_file = os.path.join(output_dir, mtrack_id + '_stem.npy')
np.save(stem_file, stem_content)
print(f'💿 Created {stem_file}')
# no preprocessing for mix audio
"""MIDI: pre-process and get metadata from the multitrack"""
notes, note_events = create_note_event_and_note_from_mtrack_mirdata(
ds, mtrack_id, fix_bass_octave=fix_bass_octave)
# save the note events and metadata to disk
notes_file = os.path.join(output_dir, mtrack_id + '_notes.npy')
np.save(notes_file, notes, allow_pickle=True, \
fix_imports=False)
print(f'🎹 Created {notes_file}')
note_events_file = os.path.join(output_dir, mtrack_id + '_note_events.npy')
np.save(note_events_file, note_events, allow_pickle=True, \
fix_imports=False)
print(f'🎹 Created {note_events_file}')
# add to the file list of the split
file_list[i] = {
'mtrack_id': mtrack_id,
'n_frames': stem_content['n_frames'], # n of audio frames
'stem_file': stem_file,
'mix_audio_file': mtrack.mix_path,
'notes_file': notes_file,
'note_events_file': note_events_file,\
'midi_file': mtrack.midi_path
}
# By split, save a file list as json
summary_dir = os.path.join(data_home, 'yourmt3_indexes')
os.makedirs(summary_dir, exist_ok=True)
summary_file = os.path.join(summary_dir, f'slakh_{split}_file_list.json')
with open(summary_file, 'w') as f:
json.dump(file_list, f, indent=4)
print(f'💾 Created {summary_file}')
elapsed_time = time.time() - start_time
print(
f"⏰: {int(elapsed_time // 3600):02d}h {int(elapsed_time % 3600 // 60):02d}m {elapsed_time % 60:.2f}s"
)
""" end of preprocess_slakh16k """
def add_program_and_is_drum_info_to_file_list(data_home: str):
for split in ['train', 'validation', 'test']:
file_list_dir = os.path.join(data_home, 'yourmt3_indexes')
file = os.path.join(file_list_dir, f'slakh_{split}_file_list.json')
with open(file, 'r') as f:
file_list = json.load(f)
for v in file_list.values():
stem_file = v['stem_file']
stem_content = np.load(stem_file, allow_pickle=True).item()
v['program'] = stem_content['program'].tolist()
v['is_drum'] = stem_content['is_drum'].tolist()
with open(file, 'w') as f:
json.dump(file_list, f, indent=4)
print(f'💾 Added program and drum info to {file}')
if __name__ == '__main__':
from config.config import shared_cfg
data_home = shared_cfg['PATH']['data_home']
preprocess_slakh16k(data_home=data_home, delete_source_files=False)