Spaces:
Running
on
T4
Running
on
T4
""" | |
Data Loaders for | |
1. contrastive learning of audio effects | |
2. music mixing style transfer | |
introduced in "Music Mixing Style Transfer: A Contrastive Learning Approach to Disentangle Audio Effects" | |
""" | |
import numpy as np | |
import wave | |
import soundfile as sf | |
import time | |
import random | |
from glob import glob | |
import torch | |
import torch.utils.data as data | |
from torch.utils.data import DataLoader | |
from torch.utils.data import Dataset | |
import os | |
import sys | |
currentdir = os.path.dirname(os.path.realpath(__file__)) | |
sys.path.append(currentdir) | |
sys.path.append(os.path.dirname(currentdir)) | |
sys.path.append(os.path.dirname(os.path.dirname(currentdir))) | |
from loader_utils import * | |
from mixing_manipulator import * | |
''' | |
Collate Functions | |
''' | |
class Collate_Variable_Length_Segments: | |
def __init__(self, args): | |
self.segment_length = args.segment_length | |
self.random_length = args.reference_length | |
self.num_strong_negatives = args.num_strong_negatives | |
if 'musdb' in args.using_dataset.lower(): | |
self.instruments = ["drums", "bass", "other", "vocals"] | |
else: | |
raise NotImplementedError | |
# collate function to trim segments A and B to random duration | |
# this function can handle different number of 'strong negative' inputs | |
def random_duration_segments_strong_negatives(self, batch): | |
num_inst = len(self.instruments) | |
# randomize current input length | |
max_length = batch[0][0].shape[-1] | |
min_length = max_length//2 | |
input_length_a, input_length_b = torch.randint(low=min_length, high=max_length, size=(2,)) | |
output_dict_A = {} | |
output_dict_B = {} | |
for cur_inst in self.instruments: | |
output_dict_A[cur_inst] = [] | |
output_dict_B[cur_inst] = [] | |
for cur_item in batch: | |
# set starting points | |
start_point_a = torch.randint(low=0, high=max_length-input_length_a, size=(1,))[0] | |
start_point_b = torch.randint(low=0, high=max_length-input_length_b, size=(1,))[0] | |
# append to output dictionary | |
for cur_i, cur_inst in enumerate(self.instruments): | |
# append A# and B# with its strong negative samples | |
for cur_neg_idx in range(self.num_strong_negatives+1): | |
output_dict_A[cur_inst].append(cur_item[cur_i*(self.num_strong_negatives+1)*2+2*cur_neg_idx][:, start_point_a : start_point_a+input_length_a]) | |
output_dict_B[cur_inst].append(cur_item[cur_i*(self.num_strong_negatives+1)*2+1+2*cur_neg_idx][:, start_point_b : start_point_b+input_length_b]) | |
''' | |
Output format : | |
[drums_A, bass_A, other_A, vocals_A], | |
[drums_B, bass_B, other_B, vocals_B] | |
''' | |
return [torch.stack(cur_segments, dim=0) for cur_inst, cur_segments in output_dict_A.items()], \ | |
[torch.stack(cur_segments, dim=0) for cur_inst, cur_segments in output_dict_B.items()] | |
# collate function for training mixing style transfer | |
def style_transfer_collate(self, batch): | |
output_dict_A1 = {} | |
output_dict_A2 = {} | |
output_dict_B2 = {} | |
for cur_inst in self.instruments: | |
output_dict_A1[cur_inst] = [] | |
output_dict_A2[cur_inst] = [] | |
output_dict_B2[cur_inst] = [] | |
for cur_item in batch: | |
# append to output dictionary | |
for cur_i, cur_inst in enumerate(self.instruments): | |
output_dict_A1[cur_inst].append(cur_item[cur_i*3]) | |
output_dict_A2[cur_inst].append(cur_item[cur_i*3+1]) | |
output_dict_B2[cur_inst].append(cur_item[cur_i*3+2]) | |
''' | |
Output format : | |
[drums_A1, bass_A1, other_A1, vocals_A1], | |
[drums_A2, bass_A2, other_A2, vocals_A2], | |
[drums_B2, bass_B2, other_B2, vocals_B2] | |
''' | |
return [torch.stack(cur_segments, dim=0) for cur_inst, cur_segments in output_dict_A1.items()], \ | |
[torch.stack(cur_segments, dim=0) for cur_inst, cur_segments in output_dict_A2.items()], \ | |
[torch.stack(cur_segments, dim=0) for cur_inst, cur_segments in output_dict_B2.items()] | |
''' | |
Data Loaders | |
''' | |
# Data loader for training the 'FXencoder' | |
# randomly loads two segments (A and B) from the dataset | |
# both segments are manipulated via FXmanipulator using (1+number of strong negative samples) sets of parameters (resulting A1, A2, ..., A#, and B1, B2, ..., B#) (# = number of strong negative samples) | |
# segments with the same effects applied (A1 and B1) are assigned as the positive pair during the training | |
# segments with the same content but with different effects applied (A2, A3, ..., A3 for A1) are also formed in a batch as 'strong negative' samples | |
# in the paper, we use strong negative samples = 1 | |
class MUSDB_Dataset_Mixing_Manipulated_FXencoder(Dataset): | |
def __init__(self, args, \ | |
mode, \ | |
applying_effects='full', \ | |
apply_prob_dict=None): | |
self.args = args | |
self.data_dir = args.data_dir + mode + "/" | |
self.mode = mode | |
self.applying_effects = applying_effects | |
self.normalization_order = args.normalization_order | |
self.fixed_random_seed = args.random_seed | |
self.pad_b4_manipulation = args.pad_b4_manipulation | |
self.pad_length = 2048 | |
if 'musdb' in args.using_dataset.lower(): | |
self.instruments = ["drums", "bass", "other", "vocals"] | |
else: | |
raise NotImplementedError | |
# path to contents | |
self.data_paths = {} | |
self.data_length_ratio_list = {} | |
# load data paths for each instrument | |
for cur_inst in self.instruments: | |
self.data_paths[cur_inst] = glob(f'{self.data_dir}{cur_inst}_normalized_{self.normalization_order}_silence_trimmed*.wav') \ | |
if args.use_normalized else glob(f'{self.data_dir}{cur_inst}_silence_trimmed*.wav') | |
self.data_length_ratio_list[cur_inst] = [] | |
# compute audio duration and its ratio | |
for cur_file_path in self.data_paths[cur_inst]: | |
cur_wav_length = load_wav_length(cur_file_path) | |
cur_inst_length_ratio = cur_wav_length / get_total_audio_length(self.data_paths[cur_inst]) | |
self.data_length_ratio_list[cur_inst].append(cur_inst_length_ratio) | |
# load effects chain | |
if applying_effects=='full': | |
if apply_prob_dict==None: | |
# initial (default) applying probabilities of each FX | |
apply_prob_dict = {'eq' : 0.9, \ | |
'comp' : 0.9, \ | |
'pan' : 0.3, \ | |
'imager' : 0.8, \ | |
'gain': 0.5} | |
reverb_prob = {'drums' : 0.5, \ | |
'bass' : 0.01, \ | |
'vocals' : 0.9, \ | |
'other' : 0.7} | |
self.mixing_manipulator = {} | |
for cur_inst in self.data_paths.keys(): | |
if 'reverb' in apply_prob_dict.keys(): | |
if cur_inst=='drums': | |
cur_reverb_weight = 0.5 | |
elif cur_inst=='bass': | |
cur_reverb_weight = 0.1 | |
else: | |
cur_reverb_weight = 1.0 | |
apply_prob_dict['reverb'] *= cur_reverb_weight | |
else: | |
apply_prob_dict['reverb'] = reverb_prob[cur_inst] | |
# create FXmanipulator for current instrument | |
self.mixing_manipulator[cur_inst] = create_inst_effects_augmentation_chain_(cur_inst, \ | |
apply_prob_dict=apply_prob_dict, \ | |
ir_dir_path=args.ir_dir_path, \ | |
sample_rate=args.sample_rate) | |
# for single effects | |
else: | |
self.mixing_manipulator = {} | |
if not isinstance(applying_effects, list): | |
applying_effects = [applying_effects] | |
for cur_inst in self.data_paths.keys(): | |
self.mixing_manipulator[cur_inst] = create_effects_augmentation_chain(applying_effects, \ | |
ir_dir_path=args.ir_dir_path) | |
def __len__(self): | |
if self.mode=='train': | |
return self.args.batch_size_total * 40 | |
else: | |
return self.args.batch_size_total | |
def __getitem__(self, idx): | |
if self.mode=="train": | |
torch.manual_seed(int(time.time())*(idx+1) % (2**32-1)) | |
np.random.seed(int(time.time())*(idx+1) % (2**32-1)) | |
random.seed(int(time.time())*(idx+1) % (2**32-1)) | |
else: | |
# fixed random seed for evaluation | |
torch.manual_seed(idx*self.fixed_random_seed) | |
np.random.seed(idx*self.fixed_random_seed) | |
random.seed(idx*self.fixed_random_seed) | |
manipulated_segments = {} | |
for cur_neg_idx in range(self.args.num_strong_negatives+1): | |
manipulated_segments[cur_neg_idx] = {} | |
# load already-saved data to save time for on-the-fly manipulation | |
cur_data_dir_path = f"{self.data_dir}manipulated_encoder/{self.args.data_save_name}/{self.applying_effects}/{idx}/" | |
if self.mode=="val" and os.path.exists(cur_data_dir_path): | |
for cur_inst in self.instruments: | |
for cur_neg_idx in range(self.args.num_strong_negatives+1): | |
cur_A_file_path = f"{cur_data_dir_path}{cur_inst}_A{cur_neg_idx+1}.wav" | |
cur_B_file_path = f"{cur_data_dir_path}{cur_inst}_B{cur_neg_idx+1}.wav" | |
cur_A = load_wav_segment(cur_A_file_path, axis=0, sample_rate=self.args.sample_rate) | |
cur_B = load_wav_segment(cur_B_file_path, axis=0, sample_rate=self.args.sample_rate) | |
manipulated_segments[cur_neg_idx][cur_inst] = [torch.from_numpy(cur_A).float(), torch.from_numpy(cur_B).float()] | |
else: | |
# repeat for number of instruments | |
for cur_inst, cur_paths in self.data_paths.items(): | |
# choose file_path to be loaded | |
cur_chosen_paths = np.random.choice(cur_paths, 2, p = self.data_length_ratio_list[cur_inst]) | |
# get random 2 starting points for each instrument | |
last_point_A = load_wav_length(cur_chosen_paths[0])-self.args.segment_length_ref | |
last_point_B = load_wav_length(cur_chosen_paths[1])-self.args.segment_length_ref | |
# simply load more data to prevent artifacts likely to be caused by the manipulator | |
if self.pad_b4_manipulation: | |
last_point_A -= self.pad_length*2 | |
last_point_B -= self.pad_length*2 | |
cur_inst_start_point_A = torch.randint(low=0, \ | |
high=last_point_A, \ | |
size=(1,))[0] | |
cur_inst_start_point_B = torch.randint(low=0, \ | |
high=last_point_B, \ | |
size=(1,))[0] | |
# load wav segments from the selected starting points | |
load_duration = self.args.segment_length_ref+self.pad_length*2 if self.pad_b4_manipulation else self.args.segment_length_ref | |
cur_inst_segment_A = load_wav_segment(cur_chosen_paths[0], \ | |
start_point=cur_inst_start_point_A, \ | |
duration=load_duration, \ | |
axis=1, \ | |
sample_rate=self.args.sample_rate) | |
cur_inst_segment_B = load_wav_segment(cur_chosen_paths[1], \ | |
start_point=cur_inst_start_point_B, \ | |
duration=load_duration, \ | |
axis=1, \ | |
sample_rate=self.args.sample_rate) | |
# mixing manipulation | |
# append A# and B# with its strong negative samples | |
for cur_neg_idx in range(self.args.num_strong_negatives+1): | |
cur_manipulated_segment_A, cur_manipulated_segment_B = self.mixing_manipulator[cur_inst]([cur_inst_segment_A, cur_inst_segment_B]) | |
# remove over-loaded area | |
if self.pad_b4_manipulation: | |
cur_manipulated_segment_A = cur_manipulated_segment_A[self.pad_length:-self.pad_length] | |
cur_manipulated_segment_B = cur_manipulated_segment_B[self.pad_length:-self.pad_length] | |
manipulated_segments[cur_neg_idx][cur_inst] = [torch.clamp(torch.transpose(torch.from_numpy(cur_manipulated_segment_A).float(), 1, 0), min=-1, max=1), \ | |
torch.clamp(torch.transpose(torch.from_numpy(cur_manipulated_segment_B).float(), 1, 0), min=-1, max=1)] | |
# check manipulated data by saving them | |
if self.mode=="val" and not os.path.exists(cur_data_dir_path): | |
os.makedirs(cur_dir_path, exist_ok=True) | |
for cur_inst in manipulated_segments[0].keys(): | |
for cur_manipulated_key, cur_manipualted_dict in manipulated_segments.items(): | |
sf.write(f"{cur_dir_path}{cur_inst}_A{cur_manipulated_key+1}.wav", torch.transpose(cur_manipualted_dict[cur_inst][0], 1, 0), self.args.sample_rate, 'PCM_16') | |
sf.write(f"{cur_dir_path}{cur_inst}_B{cur_manipulated_key+1}.wav", torch.transpose(cur_manipualted_dict[cur_inst][1], 1, 0), self.args.sample_rate, 'PCM_16') | |
output_list = [] | |
output_list_param = [] | |
for cur_inst in manipulated_segments[0].keys(): | |
for cur_manipulated_key, cur_manipualted_dict in manipulated_segments.items(): | |
output_list.extend(cur_manipualted_dict[cur_inst]) | |
''' | |
Output format: | |
list of effects manipulated stems of each instrument | |
drums_A1, drums_B1, drums_A2, drums_B2, drums_A3, drums_B3, ... , | |
bass_A1, bass_B1, bass_A2, bass_B2, bass_A3, bass_B3, ... , | |
other_A1, other_B1, other_A2, other_B2, other_A3, other_B3, ... , | |
vocals_A1, vocals_B1, vocals_A2, vocals_B2, vocals_A3, vocals_B3, ... | |
each stem has the shape of (number of channels, segment duration) | |
''' | |
return output_list | |
# generate random manipulated results for evaluation | |
def generate_contents_w_effects(self, num_content, num_effects, out_dir): | |
print(f"start generating random effects of {self.applying_effects} applied contents") | |
os.makedirs(out_dir, exist_ok=True) | |
manipulated_segments = {} | |
for cur_fx_idx in range(num_effects): | |
manipulated_segments[cur_fx_idx] = {} | |
# repeat for number of instruments | |
for cur_inst, cur_paths in self.data_paths.items(): | |
# choose file_path to be loaded | |
cur_path = np.random.choice(cur_paths, 1, p = self.data_length_ratio_list[cur_inst])[0] | |
print(f"\tgenerating instrument : {cur_inst}") | |
# get random 2 starting points for each instrument | |
last_point = load_wav_length(cur_path)-self.args.segment_length_ref | |
# simply load more data to prevent artifacts likely to be caused by the manipulator | |
if self.pad_b4_manipulation: | |
last_point -= self.pad_length*2 | |
cur_inst_start_points = torch.randint(low=0, \ | |
high=last_point, \ | |
size=(num_content,)) | |
# load wav segments from the selected starting points | |
cur_inst_segments = [] | |
for cur_num_content in range(num_content): | |
cur_ori_sample = load_wav_segment(cur_path, \ | |
start_point=cur_inst_start_points[cur_num_content], \ | |
duration=self.args.segment_length_ref, \ | |
axis=1, \ | |
sample_rate=self.args.sample_rate) | |
cur_inst_segments.append(cur_ori_sample) | |
sf.write(f"{out_dir}{cur_inst}_ori_{cur_num_content}.wav", cur_ori_sample, self.args.sample_rate, 'PCM_16') | |
# mixing manipulation | |
for cur_fx_idx in range(num_effects): | |
cur_manipulated_segments = self.mixing_manipulator[cur_inst](cur_inst_segments) | |
# remove over-loaded area | |
if self.pad_b4_manipulation: | |
for cur_man_idx in range(len(cur_manipulated_segments)): | |
cur_segment_trimmed = cur_manipulated_segments[cur_man_idx][self.pad_length:-self.pad_length] | |
cur_manipulated_segments[cur_man_idx] = torch.clamp(torch.transpose(torch.from_numpy(cur_segment_trimmed).float(), 1, 0), min=-1, max=1) | |
manipulated_segments[cur_fx_idx][cur_inst] = cur_manipulated_segments | |
# write generated data | |
# save each instruments | |
for cur_inst in manipulated_segments[0].keys(): | |
for cur_manipulated_key, cur_manipualted_dict in manipulated_segments.items(): | |
for cur_content_idx in range(num_content): | |
sf.write(f"{out_dir}{cur_inst}_{chr(65+cur_content_idx//26)}{chr(65+cur_content_idx%26)}{cur_manipulated_key+1}.wav", torch.transpose(cur_manipualted_dict[cur_inst][cur_content_idx], 1, 0), self.args.sample_rate, 'PCM_16') | |
# save mixture | |
for cur_manipulated_key, cur_manipualted_dict in manipulated_segments.items(): | |
for cur_content_idx in range(num_content): | |
for cur_idx, cur_inst in enumerate(manipulated_segments[0].keys()): | |
if cur_idx==0: | |
cur_mixture = cur_manipualted_dict[cur_inst][cur_content_idx] | |
else: | |
cur_mixture += cur_manipualted_dict[cur_inst][cur_content_idx] | |
sf.write(f"{out_dir}mixture_{chr(65+cur_content_idx//26)}{chr(65+cur_content_idx%26)}{cur_manipulated_key+1}.wav", torch.transpose(cur_mixture, 1, 0), self.args.sample_rate, 'PCM_16') | |
return | |
# Data loader for training the 'Mastering Style Converter' | |
# loads two segments (A and B) from the dataset | |
# both segments are manipulated via Mastering Effects Manipulator (resulting A1, A2, and B2) | |
# one of the manipulated segment is used as a reference segment (B2), which is randomly manipulated the same as the ground truth segment (A2) | |
class MUSDB_Dataset_Mixing_Manipulated_Style_Transfer(Dataset): | |
def __init__(self, args, \ | |
mode, \ | |
applying_effects='full', \ | |
apply_prob_dict=None): | |
self.args = args | |
self.data_dir = args.data_dir + mode + "/" | |
self.mode = mode | |
self.applying_effects = applying_effects | |
self.fixed_random_seed = args.random_seed | |
self.pad_b4_manipulation = args.pad_b4_manipulation | |
self.pad_length = 2048 | |
if 'musdb' in args.using_dataset.lower(): | |
self.instruments = ["drums", "bass", "other", "vocals"] | |
else: | |
raise NotImplementedError | |
# load data paths for each instrument | |
self.data_paths = {} | |
self.data_length_ratio_list = {} | |
for cur_inst in self.instruments: | |
self.data_paths[cur_inst] = glob(f'{self.data_dir}{cur_inst}_normalized_{self.args.normalization_order}_silence_trimmed*.wav') \ | |
if args.use_normalized else glob(f'{self.data_dir}{cur_inst}_silence_trimmed.wav') | |
self.data_length_ratio_list[cur_inst] = [] | |
# compute audio duration and its ratio | |
for cur_file_path in self.data_paths[cur_inst]: | |
cur_wav_length = load_wav_length(cur_file_path) | |
cur_inst_length_ratio = cur_wav_length / get_total_audio_length(self.data_paths[cur_inst]) | |
self.data_length_ratio_list[cur_inst].append(cur_inst_length_ratio) | |
self.mixing_manipulator = {} | |
if applying_effects=='full': | |
if apply_prob_dict==None: | |
# initial (default) applying probabilities of each FX | |
# we don't update these probabilities for training the MixFXcloner | |
apply_prob_dict = {'eq' : 0.9, \ | |
'comp' : 0.9, \ | |
'pan' : 0.3, \ | |
'imager' : 0.8, \ | |
'gain': 0.5} | |
reverb_prob = {'drums' : 0.5, \ | |
'bass' : 0.01, \ | |
'vocals' : 0.9, \ | |
'other' : 0.7} | |
for cur_inst in self.data_paths.keys(): | |
if 'reverb' in apply_prob_dict.keys(): | |
if cur_inst=='drums': | |
cur_reverb_weight = 0.5 | |
elif cur_inst=='bass': | |
cur_reverb_weight = 0.1 | |
else: | |
cur_reverb_weight = 1.0 | |
apply_prob_dict['reverb'] *= cur_reverb_weight | |
else: | |
apply_prob_dict['reverb'] = reverb_prob[cur_inst] | |
self.mixing_manipulator[cur_inst] = create_inst_effects_augmentation_chain(cur_inst, \ | |
apply_prob_dict=apply_prob_dict, \ | |
ir_dir_path=args.ir_dir_path, \ | |
sample_rate=args.sample_rate) | |
# for single effects | |
else: | |
if not isinstance(applying_effects, list): | |
applying_effects = [applying_effects] | |
for cur_inst in self.data_paths.keys(): | |
self.mixing_manipulator[cur_inst] = create_effects_augmentation_chain(applying_effects, \ | |
ir_dir_path=args.ir_dir_path) | |
def __len__(self): | |
min_length = get_total_audio_length(glob(f'{self.data_dir}vocals_normalized_{self.args.normalization_order}*.wav')) | |
data_len = min_length // self.args.segment_length | |
return data_len | |
def __getitem__(self, idx): | |
if self.mode=="train": | |
torch.manual_seed(int(time.time())*(idx+1) % (2**32-1)) | |
np.random.seed(int(time.time())*(idx+1) % (2**32-1)) | |
random.seed(int(time.time())*(idx+1) % (2**32-1)) | |
else: | |
# fixed random seed for evaluation | |
torch.manual_seed(idx*self.fixed_random_seed) | |
np.random.seed(idx*self.fixed_random_seed) | |
random.seed(idx*self.fixed_random_seed) | |
manipulated_segments = {} | |
# load already-saved data to save time for on-the-fly manipulation | |
cur_data_dir_path = f"{self.data_dir}manipulated_converter/{self.args.data_save_name}/{self.applying_effects}/{idx}/" | |
if self.mode=="val" and os.path.exists(cur_data_dir_path): | |
for cur_inst in self.instruments: | |
cur_A1_file_path = f"{cur_data_dir_path}{cur_inst}_A1.wav" | |
cur_A2_file_path = f"{cur_data_dir_path}{cur_inst}_A2.wav" | |
cur_B2_file_path = f"{cur_data_dir_path}{cur_inst}_B2.wav" | |
cur_manipulated_segment_A1 = load_wav_segment(cur_A1_file_path, axis=0, sample_rate=self.args.sample_rate) | |
cur_manipulated_segment_A2 = load_wav_segment(cur_A2_file_path, axis=0, sample_rate=self.args.sample_rate) | |
cur_manipulated_segment_B2 = load_wav_segment(cur_B2_file_path, axis=0, sample_rate=self.args.sample_rate) | |
manipulated_segments[cur_inst] = [torch.from_numpy(cur_manipulated_segment_A1).float(), \ | |
torch.from_numpy(cur_manipulated_segment_A2).float(), \ | |
torch.from_numpy(cur_manipulated_segment_B2).float()] | |
else: | |
# repeat for number of instruments | |
for cur_inst, cur_paths in self.data_paths.items(): | |
# choose file_path to be loaded | |
cur_chosen_paths = np.random.choice(cur_paths, 2, p = self.data_length_ratio_list[cur_inst]) | |
# cur_chosen_paths = [cur_paths[idx], cur_paths[idx+1]] | |
# get random 2 starting points for each instrument | |
last_point_A = load_wav_length(cur_chosen_paths[0])-self.args.segment_length_ref | |
last_point_B = load_wav_length(cur_chosen_paths[1])-self.args.segment_length_ref | |
# simply load more data to prevent artifacts likely to be caused by the manipulator | |
if self.pad_b4_manipulation: | |
last_point_A -= self.pad_length*2 | |
last_point_B -= self.pad_length*2 | |
cur_inst_start_point_A = torch.randint(low=0, \ | |
high=last_point_A, \ | |
size=(1,))[0] | |
cur_inst_start_point_B = torch.randint(low=0, \ | |
high=last_point_B, \ | |
size=(1,))[0] | |
# load wav segments from the selected starting points | |
load_duration = self.args.segment_length_ref+self.pad_length*2 if self.pad_b4_manipulation else self.args.segment_length_ref | |
cur_inst_segment_A = load_wav_segment(cur_chosen_paths[0], \ | |
start_point=cur_inst_start_point_A, \ | |
duration=load_duration, \ | |
axis=1, \ | |
sample_rate=self.args.sample_rate) | |
cur_inst_segment_B = load_wav_segment(cur_chosen_paths[1], \ | |
start_point=cur_inst_start_point_B, \ | |
duration=load_duration, \ | |
axis=1, \ | |
sample_rate=self.args.sample_rate) | |
''' mixing manipulation ''' | |
# manipulate segment A and B to produce | |
# input : A1 (normalized sample) | |
# ground truth : A2 | |
# reference : B2 | |
cur_manipulated_segment_A1 = cur_inst_segment_A | |
cur_manipulated_segment_A2, cur_manipulated_segment_B2 = self.mixing_manipulator[cur_inst]([cur_inst_segment_A, cur_inst_segment_B]) | |
# remove over-loaded area | |
if self.pad_b4_manipulation: | |
cur_manipulated_segment_A1 = cur_manipulated_segment_A1[self.pad_length:-self.pad_length] | |
cur_manipulated_segment_A2 = cur_manipulated_segment_A2[self.pad_length:-self.pad_length] | |
cur_manipulated_segment_B2 = cur_manipulated_segment_B2[self.pad_length:-self.pad_length] | |
manipulated_segments[cur_inst] = [torch.clamp(torch.transpose(torch.from_numpy(cur_manipulated_segment_A1).float(), 1, 0), min=-1, max=1), \ | |
torch.clamp(torch.transpose(torch.from_numpy(cur_manipulated_segment_A2).float(), 1, 0), min=-1, max=1), \ | |
torch.clamp(torch.transpose(torch.from_numpy(cur_manipulated_segment_B2).float(), 1, 0), min=-1, max=1)] | |
# check manipulated data by saving them | |
if (self.mode=="val" and not os.path.exists(cur_data_dir_path)): | |
mixture_dict = {} | |
for cur_inst in manipulated_segments.keys(): | |
cur_inst_dir_path = f"{cur_data_dir_path}{idx}/{cur_inst}/" | |
os.makedirs(cur_inst_dir_path, exist_ok=True) | |
sf.write(f"{cur_inst_dir_path}A1.wav", torch.transpose(manipulated_segments[cur_inst][0], 1, 0), self.args.sample_rate, 'PCM_16') | |
sf.write(f"{cur_inst_dir_path}A2.wav", torch.transpose(manipulated_segments[cur_inst][1], 1, 0), self.args.sample_rate, 'PCM_16') | |
sf.write(f"{cur_inst_dir_path}B2.wav", torch.transpose(manipulated_segments[cur_inst][2], 1, 0), self.args.sample_rate, 'PCM_16') | |
mixture_dict['A1'] = torch.transpose(manipulated_segments[cur_inst][0], 1, 0) | |
mixture_dict['A2'] = torch.transpose(manipulated_segments[cur_inst][1], 1, 0) | |
mixture_dict['B2'] = torch.transpose(manipulated_segments[cur_inst][2], 1, 0) | |
cur_mix_dir_path = f"{cur_data_dir_path}{idx}/mixture/" | |
os.makedirs(cur_mix_dir_path, exist_ok=True) | |
sf.write(f"{cur_mix_dir_path}A1.wav", mixture_dict['A1'], self.args.sample_rate, 'PCM_16') | |
sf.write(f"{cur_mix_dir_path}A2.wav", mixture_dict['A2'], self.args.sample_rate, 'PCM_16') | |
sf.write(f"{cur_mix_dir_path}B2.wav", mixture_dict['B2'], self.args.sample_rate, 'PCM_16') | |
output_list = [] | |
for cur_inst in manipulated_segments.keys(): | |
output_list.extend(manipulated_segments[cur_inst]) | |
''' | |
Output format: | |
list of effects manipulated stems of each instrument | |
drums_A1, drums_A2, drums_B2, | |
bass_A1, bass_A2, bass_B2, | |
other_A1, other_A2, other_B2, | |
vocals_A1, vocals_A2, vocals_B2, | |
each stem has the shape of (number of channels, segment duration) | |
Notation : | |
A1 = input of the network | |
A2 = ground truth | |
B2 = reference track | |
''' | |
return output_list | |
# Data loader for inferencing the task 'Mixing Style Transfer' | |
### loads whole mixture or stems from the target directory | |
class Song_Dataset_Inference(Dataset): | |
def __init__(self, args): | |
self.args = args | |
self.data_dir = args.target_dir | |
self.interpolate = args.interpolation | |
self.instruments = args.instruments | |
self.data_dir_paths = sorted(glob(f"{self.data_dir}*/")) | |
self.input_name = args.input_file_name | |
self.reference_name = args.reference_file_name | |
self.stem_level_directory_name = args.stem_level_directory_name \ | |
if self.args.do_not_separate else os.path.join(args.stem_level_directory_name, args.separation_model) | |
if self.interpolate: | |
self.reference_name_B = args.reference_file_name_2interpolate | |
# audio effects normalizer | |
if args.normalize_input: | |
self.normalization_chain = Audio_Effects_Normalizer(precomputed_feature_path=args.precomputed_normalization_feature, \ | |
STEMS=args.instruments, \ | |
EFFECTS=args.normalization_order) | |
def __len__(self): | |
return len(self.data_dir_paths) | |
def __getitem__(self, idx): | |
''' stem-level conversion ''' | |
input_stems = [] | |
reference_stems = [] | |
reference_B_stems = [] | |
for cur_inst in self.instruments: | |
cur_input_file_path = os.path.join(self.data_dir_paths[idx], self.stem_level_directory_name, self.input_name, cur_inst+'.wav') | |
cur_reference_file_path = os.path.join(self.data_dir_paths[idx], self.stem_level_directory_name, self.reference_name, cur_inst+'.wav') | |
# load wav | |
cur_input_wav = load_wav_segment(cur_input_file_path, axis=0, sample_rate=self.args.sample_rate) | |
cur_reference_wav = load_wav_segment(cur_reference_file_path, axis=0, sample_rate=self.args.sample_rate) | |
if self.args.normalize_input: | |
cur_input_wav = self.normalization_chain.normalize_audio(cur_input_wav.transpose(), src=cur_inst).transpose() | |
input_stems.append(torch.clamp(torch.from_numpy(cur_input_wav).float(), min=-1, max=1)) | |
reference_stems.append(torch.clamp(torch.from_numpy(cur_reference_wav).float(), min=-1, max=1)) | |
# for interpolation | |
if self.interpolate: | |
cur_reference_B_file_path = os.path.join(self.data_dir_paths[idx], self.stem_level_directory_name, self.reference_name_B, cur_inst+'.wav') | |
cur_reference_B_wav = load_wav_segment(cur_reference_B_file_path, axis=0, sample_rate=self.args.sample_rate) | |
reference_B_stems.append(torch.clamp(torch.from_numpy(cur_reference_B_wav).float(), min=-1, max=1)) | |
dir_name = os.path.dirname(self.data_dir_paths[idx]) | |
if self.interpolate: | |
return torch.stack(input_stems, dim=0), torch.stack(reference_stems, dim=0), torch.stack(reference_B_stems, dim=0), dir_name | |
else: | |
return torch.stack(input_stems, dim=0), torch.stack(reference_stems, dim=0), dir_name | |
# check dataset | |
if __name__ == '__main__': | |
""" | |
Test code of data loaders | |
""" | |
import time | |
print('checking dataset...') | |
total_epochs = 1 | |
bs = 5 | |
check_step_size = 3 | |
collate_class = Collate_Variable_Length_Segments(args) | |
print('\n========== Effects Encoder ==========') | |
from config import args | |
##### generate samples with ranfom configuration | |
# args.normalization_order = 'eqcompimagegain' | |
# for cur_effect in ['full', 'gain', 'comp', 'reverb', 'eq', 'imager', 'pan']: | |
# start_time = time.time() | |
# dataset = MUSDB_Dataset_Mixing_Manipulated_FXencoder(args, mode='val', applying_effects=cur_effect, check_data=True) | |
# dataset.generate_contents_w_effects(num_content=25, num_effects=10) | |
# print(f'\t---time taken : {time.time()-start_time}---') | |
### training data loder | |
dataset = MUSDB_Dataset_Mixing_Manipulated_FXencoder(args, mode='train', applying_effects=['comp']) | |
data_loader = DataLoader(dataset, \ | |
batch_size=bs, \ | |
shuffle=False, \ | |
collate_fn=collate_class.random_duration_segments_strong_negatives, \ | |
drop_last=False, \ | |
num_workers=0) | |
for epoch in range(total_epochs): | |
start_time_loader = time.time() | |
for step, output_list in enumerate(data_loader): | |
if step==check_step_size: | |
break | |
print(f'Epoch {epoch+1}/{total_epochs}\tStep {step+1}/{len(data_loader)}') | |
print(f'num contents : {len(output_list)}\tnum instruments : {len(output_list[0])}\tcontent A shape : {output_list[0][0].shape}\t content B shape : {output_list[1][0].shape} \ttime taken: {time.time()-start_time_loader:.4f}') | |
start_time_loader = time.time() | |
print('\n========== Mixing Style Transfer ==========') | |
from trainer_mixing_transfer.config_conv import args | |
### training data loder | |
dataset = MUSDB_Dataset_Mixing_Manipulated_Style_Transfer(args, mode='train') | |
data_loader = DataLoader(dataset, \ | |
batch_size=bs, \ | |
shuffle=False, \ | |
collate_fn=collate_class.style_transfer_collate, \ | |
drop_last=False, \ | |
num_workers=0) | |
for epoch in range(total_epochs): | |
start_time_loader = time.time() | |
for step, output_list in enumerate(data_loader): | |
if step==check_step_size: | |
break | |
print(f'Epoch {epoch+1}/{total_epochs}\tStep {step+1}/{len(data_loader)}') | |
print(f'num contents : {len(output_list)}\tnum instruments : {len(output_list[0])}\tA1 shape : {output_list[0][0].shape}\tA2 shape : {output_list[1][0].shape}\tA3 shape : {output_list[2][0].shape}\ttime taken: {time.time()-start_time_loader:.4f}') | |
start_time_loader = time.time() | |
print('\n--- checking dataset completed ---') | |