Making coherent music mixtures with Scaper

make_coherent_mix.py

import sys
sys.path.insert(0, '.')

from multiprocessing import cpu_count
import os
from utils import parallel_process
from scaper import Scaper, generate_from_jams
import librosa
import numpy as np
from tqdm import tqdm
import sox
import os

# Assuming data @ "/storage/data/musdb/train"]
# Assuming data @ "/storage/data/musdb/validation"]


mixture_parameters = {
    'training': {
        'num_mixtures': 20000,
        'foreground_path': '/storage/data/musdb/train',
        'background_path': '/storage/data/musdb/train',
        'target_path': '/storage/data/musdb/generated/train',
        'scene_duration': 10
    },
    'validation': {
        'num_mixtures': 1000,
        'foreground_path': '/storage/data/musdb/validation',
        'background_path': '/storage/data/musdb/validation',
        'target_path': '/storage/data/musdb/generated/validation',
        'scene_duration': 10
    },
}

silent_files = {'default': True}

def instantiate_and_get_event_spec(sc, master_label, scene_duration):
    file_name = 'default'
    while silent_files[file_name]:
        sc.add_event(
                label=('const', master_label),
                source_file=('choose', []),
                source_time=('uniform', 0, 1000),
                event_time=('const', 0),
                event_duration=('const', sc.duration),
                snr=('const', 0.0),
                pitch_shift=('uniform', -2, 2),
                time_stretch=None
        )
        event = sc._instantiate_event(sc.fg_spec[0])
        if event.source_file not in silent_files:
            silent_files[event.source_file] = sox.file_info.silent(event.source_file)
        file_name = event.source_file
        
    if event.source_file not in timings:
        audio, sr = librosa.load(event.source_file, sr=None, mono=True)
        duration = audio.shape[0] / float(sr)
        rms = librosa.feature.rmse(audio)[0, :]
        loud = np.argwhere(rms > np.percentile(rms, 50))[:, 0]
        times = librosa.frames_to_time(loud, sr=sr).tolist()
        timings[event.source_file] = times
    else:
        times = timings[event.source_file]

    sc.fg_spec = []
    sc.add_event(
            label=('const', 'vocals'),
            source_file=('choose', []),
            source_time=('choose', times),
            event_time=('const', 0),
            event_duration=('const', sc.duration),
            snr=('uniform', 0, 5),
            pitch_shift=('uniform', -2, 2),
            time_stretch=None
    )
    event = sc._instantiate_event(sc.fg_spec[0])
    sc.fg_spec = []
    return sc, event

def make_one_mixture(sc, path_to_file, labels):
    sc, event = instantiate_and_get_event_spec(sc, labels[0], sc.duration)
    for label in labels:
        sc.add_event(
            label=('const', label),
            source_file=('const', event.source_file.replace('vocals', label)),
            source_time=('const', event.source_time),
            event_time=('const', 0),
            event_duration=('const', sc.duration),
            snr=('uniform', 0, 5),
            pitch_shift=('const', event.pitch_shift),
            time_stretch=None
        )
    sc.generate(path_to_file, path_to_file.replace('.wav', '.jams'), no_audio=True)
    sc.fg_spec = []


for key, params in mixture_parameters.items():
    labels = ['vocals', 'bass', 'drums', 'other']
    sc = Scaper(
        params['scene_duration'], 
        fg_path=params['foreground_path'], 
        bg_path=params['foreground_path'], 
        random_state=0
    )
    os.makedirs(params['target_path'], exist_ok=True)
    sc.ref_db = -25
    timings = {}

    for i in tqdm(range(params['num_mixtures'])):
        path_to_file = os.path.join(params['target_path'], f'{i:08d}.wav')
        make_one_mixture(sc, path_to_file, labels)

def synthesize_one_mixture(jams_file):
    wav_file = jams_file.replace('.jams', '.wav')
    generate_from_jams(jams_file, wav_file, save_sources=True)

def synthesize_mixtures_in_parallel(target_path, n_jobs):
    jam_files = [
        os.path.join(target_path, x) for x in os.listdir(target_path) if '.jams' in x
    ]
    parallel_process(
        jam_files, synthesize_one_mixture, cpu_count()
    )

for key, params in mixture_parameters.items():
    synthesize_mixtures_in_parallel(params['target_path'], cpu_count())