Generating coherent music with Scaper

coherent_music_generation_with_scaper.py

import scaper
import os
import sox
import numpy as np
import logging
import warnings
import random
import jams
import librosa
import itertools

warnings.filterwarnings("ignore", category=scaper.scaper_warnings.ScaperWarning)
logging.getLogger('scaper').setLevel(logging.ERROR)
warnings.simplefilter(action='ignore', category=FutureWarning)
logging.getLogger('sox').setLevel(logging.ERROR)

def create_scaper_generator(foreground_directory, scene_duration, song_title, song_start, background_directory=None):
    if background_directory is None:
        background_directory = foreground_directory
    sc = scaper.Scaper(scene_duration, 
                       foreground_directory, 
                       background_directory)
    sc.protected_labels = []
    sc.ref_db = -50
    sc.n_channels = 1
    sc.sr = 44100
    sc.min_silence_duration = None

    n_events = 4 

    labels = ['vocals', 'drums', 'bass', 'other']

    for n in range(n_events):
        sc.add_event(label=('const', labels[n]),
                    source_file=('const', os.path.join(foreground_directory, labels[n], song_title + '.wav')),
                    source_time=('const', song_start),
                    event_time=('const', 0),
                    event_duration=('const', scene_duration),
                    snr=('uniform', 25, 30),
                    pitch_shift=None,
                    time_stretch=None)
    return sc
  
def check_sources_not_equal_to_mix(jams_path):
    mix, sr = librosa.load(jams_path.replace('.json', '.wav'), sr=None)
    jam = jams.load(jams_path)
    data = jam.annotations[0]['data']['value']
    sources = [mix]
    for d in data:
        if d['role'] == 'foreground':  
            sources.append(librosa.load(d['saved_source_file'], sr=None)[0])
    if len(set([x.shape[0] for x in sources])) > 1:
        print('Not all equal length!')
        return True
    return False

def create_mixture(i, song_title, song_start, scene_duration, foreground_directory, target_directory):
    sc = create_scaper_generator(foreground_directory, scene_duration, song_title, song_start)
    audio_path = os.path.join(target_directory, '%06d.wav' % i)
    jams_path = os.path.join(target_directory, '%06d.json' % i)
    error = True
    while error:
        try:
            sc.generate(audio_path, jams_path, save_sources=True, allow_repeated_label=False, reverb=0.0)
            error = check_sources_not_equal_to_mix(jams_path)
        except KeyboardInterrupt:
            raise
        except:
            pass

base_directory = '/path/to/generated-soundscapes-directory/'
if not os.path.exists(base_directory):
    os.mkdir(base_directory)

directories = {
    'training': os.path.join('/path/to/musdb-training-source-folders', 'train'),
  'validation': os.path.join('/path/to/musdb-validation-source-folders', 'validation'),
  'testing': os.path.join('/path/to/musdb-testing-source-folders', 'test'),
  'background': None
}

songs = {
  'training': [(x[:-4], sox.file_info.duration(os.path.join(directories['training'], 'mixture', x))) for x in os.listdir(os.path.join(directories['training'], 'mixture'))],
  'validation': [(x[:-4], sox.file_info.duration(os.path.join(directories['validation'], 'mixture', x))) for x in os.listdir(os.path.join(directories['validation'], 'mixture'))],
}

num_training = 20000
num_validation = 2000
n_jobs = 24
splits = [('training', num_training), ('validation', num_validation)]

for split, num_split in splits:
    target_directory = os.path.join(base_directory, split)
    if not os.path.exists(target_directory):
        os.mkdir(target_directory)
    mixes = []
    for i in range(num_split):
        song_title, song_duration = random.choice(songs[split])
        scene_duration = 3.2
        song_start_time = random.uniform(0, song_duration - scene_duration)
        mixes.append({"i": i, 
                   'song_title': song_title,
                   'song_start': song_start_time,
                   'scene_duration': 3.2,
                   'foreground_directory': directories[split],
                   "target_directory": target_directory})
    for mix in mixes: #can be parallelized
        create_mixture(**mix)