ASesYusuf1/SESA_Audio_Separation
14
1# coding: utf-82__author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/'3 4 5import os6import random7import numpy as np8import torch9import soundfile as sf10import pickle11import time12import itertools13import multiprocessing14from tqdm.auto import tqdm15from glob import glob16import audiomentations as AU17import pedalboard as PB18import warnings19warnings.filterwarnings("ignore")20 21 22def load_chunk(path, length, chunk_size, offset=None):23 if chunk_size <= length:24 if offset is None:25 offset = np.random.randint(length - chunk_size + 1)26 x = sf.read(path, dtype='float32', start=offset, frames=chunk_size)[0]27 else:28 x = sf.read(path, dtype='float32')[0]29 if len(x.shape) == 1:30 # Mono case31 pad = np.zeros((chunk_size - length))32 else:33 pad = np.zeros([chunk_size - length, x.shape[-1]])34 x = np.concatenate([x, pad], axis=0)35 # Mono fix36 if len(x.shape) == 1:37 x = np.expand_dims(x, axis=1)38 return x.T39 40 41def get_track_set_length(params):42 path, instruments, file_types = params43 # Check lengths of all instruments (it can be different in some cases)44 lengths_arr = []45 for instr in instruments:46 length = -147 for extension in file_types:48 path_to_audio_file = path + '/{}.{}'.format(instr, extension)49 if os.path.isfile(path_to_audio_file):50 length = len(sf.read(path_to_audio_file)[0])51 break52 if length == -1:53 print('Cant find file "{}" in folder {}'.format(instr, path))54 continue55 lengths_arr.append(length)56 lengths_arr = np.array(lengths_arr)57 if lengths_arr.min() != lengths_arr.max():58 print('Warning: lengths of stems are different for path: {}. ({} != {})'.format(59 path,60 lengths_arr.min(),61 lengths_arr.max())62 )63 # We use minimum to allow overflow for soundfile read in non-equal length cases64 return path, lengths_arr.min()65 66 67# For multiprocessing68def get_track_length(params):69 path = params70 length = len(sf.read(path)[0])71 return (path, length)72 73 74class MSSDataset(torch.utils.data.Dataset):75 def __init__(self, config, data_path, metadata_path="metadata.pkl", dataset_type=1, batch_size=None, verbose=True):76 self.verbose = verbose77 self.config = config78 self.dataset_type = dataset_type # 1, 2, 3 or 479 self.data_path = data_path80 self.instruments = instruments = config.training.instruments81 if batch_size is None:82 batch_size = config.training.batch_size83 self.batch_size = batch_size84 self.file_types = ['wav', 'flac']85 self.metadata_path = metadata_path86 87 # Augmentation block88 self.aug = False89 if 'augmentations' in config:90 if config['augmentations'].enable is True:91 if self.verbose:92 print('Use augmentation for training')93 self.aug = True94 else:95 if self.verbose:96 print('There is no augmentations block in config. Augmentations disabled for training...')97 98 metadata = self.get_metadata()99 100 if self.dataset_type in [1, 4]:101 if len(metadata) > 0:102 if self.verbose:103 print('Found tracks in dataset: {}'.format(len(metadata)))104 else:105 print('No tracks found for training. Check paths you provided!')106 exit()107 else:108 for instr in self.instruments:109 if self.verbose:110 print('Found tracks for {} in dataset: {}'.format(instr, len(metadata[instr])))111 self.metadata = metadata112 self.chunk_size = config.audio.chunk_size113 self.min_mean_abs = config.audio.min_mean_abs114 115 def __len__(self):116 return self.config.training.num_steps * self.batch_size117 118 def read_from_metadata_cache(self, track_paths, instr=None):119 metadata = []120 if os.path.isfile(self.metadata_path):121 if self.verbose:122 print('Found metadata cache file: {}'.format(self.metadata_path))123 old_metadata = pickle.load(open(self.metadata_path, 'rb'))124 else:125 return track_paths, metadata126 127 if instr:128 old_metadata = old_metadata[instr]129 130 # We will not re-read tracks existed in old metadata file131 track_paths_set = set(track_paths)132 for old_path, file_size in old_metadata:133 if old_path in track_paths_set:134 metadata.append([old_path, file_size])135 track_paths_set.remove(old_path)136 track_paths = list(track_paths_set)137 if len(metadata) > 0:138 print('Old metadata was used for {} tracks.'.format(len(metadata)))139 return track_paths, metadata140 141 142 def get_metadata(self):143 read_metadata_procs = multiprocessing.cpu_count()144 if 'read_metadata_procs' in self.config['training']:145 read_metadata_procs = int(self.config['training']['read_metadata_procs'])146 147 if self.verbose:148 print(149 'Dataset type:', self.dataset_type,150 'Processes to use:', read_metadata_procs,151 '\nCollecting metadata for', str(self.data_path),152 )153 154 if self.dataset_type in [1, 4]:155 track_paths = []156 if type(self.data_path) == list:157 for tp in self.data_path:158 tracks_for_folder = sorted(glob(tp + '/*'))159 if len(tracks_for_folder) == 0:160 print('Warning: no tracks found in folder \'{}\'. Please check it!'.format(tp))161 track_paths += tracks_for_folder162 else:163 track_paths += sorted(glob(self.data_path + '/*'))164 165 track_paths = [path for path in track_paths if os.path.basename(path)[0] != '.' and os.path.isdir(path)]166 track_paths, metadata = self.read_from_metadata_cache(track_paths, None)167 168 if read_metadata_procs <= 1:169 for path in tqdm(track_paths):170 track_path, track_length = get_track_set_length((path, self.instruments, self.file_types))171 metadata.append((track_path, track_length))172 else:173 p = multiprocessing.Pool(processes=read_metadata_procs)174 with tqdm(total=len(track_paths)) as pbar:175 track_iter = p.imap(176 get_track_set_length,177 zip(track_paths, itertools.repeat(self.instruments), itertools.repeat(self.file_types))178 )179 for track_path, track_length in track_iter:180 metadata.append((track_path, track_length))181 pbar.update()182 p.close()183 184 elif self.dataset_type == 2:185 metadata = dict()186 for instr in self.instruments:187 metadata[instr] = []188 track_paths = []189 if type(self.data_path) == list:190 for tp in self.data_path:191 track_paths += sorted(glob(tp + '/{}/*.wav'.format(instr)))192 track_paths += sorted(glob(tp + '/{}/*.flac'.format(instr)))193 else:194 track_paths += sorted(glob(self.data_path + '/{}/*.wav'.format(instr)))195 track_paths += sorted(glob(self.data_path + '/{}/*.flac'.format(instr)))196 197 track_paths, metadata[instr] = self.read_from_metadata_cache(track_paths, instr)198 199 if read_metadata_procs <= 1:200 for path in tqdm(track_paths):201 length = len(sf.read(path)[0])202 metadata[instr].append((path, length))203 else:204 p = multiprocessing.Pool(processes=read_metadata_procs)205 for out in tqdm(p.imap(get_track_length, track_paths), total=len(track_paths)):206 metadata[instr].append(out)207 208 elif self.dataset_type == 3:209 import pandas as pd210 if type(self.data_path) != list:211 data_path = [self.data_path]212 213 metadata = dict()214 for i in range(len(self.data_path)):215 if self.verbose:216 print('Reading tracks from: {}'.format(self.data_path[i]))217 df = pd.read_csv(self.data_path[i])218 219 skipped = 0220 for instr in self.instruments:221 part = df[df['instrum'] == instr].copy()222 print('Tracks found for {}: {}'.format(instr, len(part)))223 for instr in self.instruments:224 part = df[df['instrum'] == instr].copy()225 metadata[instr] = []226 track_paths = list(part['path'].values)227 track_paths, metadata[instr] = self.read_from_metadata_cache(track_paths, instr)228 229 for path in tqdm(track_paths):230 if not os.path.isfile(path):231 print('Cant find track: {}'.format(path))232 skipped += 1233 continue234 # print(path)235 try:236 length = len(sf.read(path)[0])237 except:238 print('Problem with path: {}'.format(path))239 skipped += 1240 continue241 metadata[instr].append((path, length))242 if skipped > 0:243 print('Missing tracks: {} from {}'.format(skipped, len(df)))244 else:245 print('Unknown dataset type: {}. Must be 1, 2, 3 or 4'.format(self.dataset_type))246 exit()247 248 # Save metadata249 pickle.dump(metadata, open(self.metadata_path, 'wb'))250 return metadata251 252 def load_source(self, metadata, instr):253 while True:254 if self.dataset_type in [1, 4]:255 track_path, track_length = random.choice(metadata)256 for extension in self.file_types:257 path_to_audio_file = track_path + '/{}.{}'.format(instr, extension)258 if os.path.isfile(path_to_audio_file):259 try:260 source = load_chunk(path_to_audio_file, track_length, self.chunk_size)261 except Exception as e:262 # Sometimes error during FLAC reading, catch it and use zero stem263 print('Error: {} Path: {}'.format(e, path_to_audio_file))264 source = np.zeros((2, self.chunk_size), dtype=np.float32)265 break266 else:267 track_path, track_length = random.choice(metadata[instr])268 try:269 source = load_chunk(track_path, track_length, self.chunk_size)270 except Exception as e:271 # Sometimes error during FLAC reading, catch it and use zero stem272 print('Error: {} Path: {}'.format(e, track_path))273 source = np.zeros((2, self.chunk_size), dtype=np.float32)274 275 if np.abs(source).mean() >= self.min_mean_abs: # remove quiet chunks276 break277 if self.aug:278 source = self.augm_data(source, instr)279 return torch.tensor(source, dtype=torch.float32)280 281 def load_random_mix(self):282 res = []283 for instr in self.instruments:284 s1 = self.load_source(self.metadata, instr)285 # Mixup augmentation. Multiple mix of same type of stems286 if self.aug:287 if 'mixup' in self.config['augmentations']:288 if self.config['augmentations'].mixup:289 mixup = [s1]290 for prob in self.config.augmentations.mixup_probs:291 if random.uniform(0, 1) < prob:292 s2 = self.load_source(self.metadata, instr)293 mixup.append(s2)294 mixup = torch.stack(mixup, dim=0)295 loud_values = np.random.uniform(296 low=self.config.augmentations.loudness_min,297 high=self.config.augmentations.loudness_max,298 size=(len(mixup),)299 )300 loud_values = torch.tensor(loud_values, dtype=torch.float32)301 mixup *= loud_values[:, None, None]302 s1 = mixup.mean(dim=0, dtype=torch.float32)303 res.append(s1)304 res = torch.stack(res)305 return res306 307 def load_aligned_data(self):308 track_path, track_length = random.choice(self.metadata)309 attempts = 10310 while attempts:311 if track_length >= self.chunk_size:312 common_offset = np.random.randint(track_length - self.chunk_size + 1)313 else:314 common_offset = None315 res = []316 silent_chunks = 0317 for i in self.instruments:318 for extension in self.file_types:319 path_to_audio_file = track_path + '/{}.{}'.format(i, extension)320 if os.path.isfile(path_to_audio_file):321 try:322 source = load_chunk(path_to_audio_file, track_length, self.chunk_size, offset=common_offset)323 except Exception as e:324 # Sometimes error during FLAC reading, catch it and use zero stem325 print('Error: {} Path: {}'.format(e, path_to_audio_file))326 source = np.zeros((2, self.chunk_size), dtype=np.float32)327 break328 res.append(source)329 if np.abs(source).mean() < self.min_mean_abs: # remove quiet chunks330 silent_chunks += 1331 if silent_chunks == 0:332 break333 334 attempts -= 1335 if attempts <= 0:336 print('Attempts max!', track_path)337 if common_offset is None:338 # If track is too small break immediately339 break340 341 res = np.stack(res, axis=0)342 if self.aug:343 for i, instr in enumerate(self.instruments):344 res[i] = self.augm_data(res[i], instr)345 return torch.tensor(res, dtype=torch.float32)346 347 def augm_data(self, source, instr):348 # source.shape = (2, 261120) - first channels, second length349 source_shape = source.shape350 applied_augs = []351 if 'all' in self.config['augmentations']:352 augs = self.config['augmentations']['all']353 else:354 augs = dict()355 356 # We need to add to all augmentations specific augs for stem. And rewrite values if needed357 if instr in self.config['augmentations']:358 for el in self.config['augmentations'][instr]:359 augs[el] = self.config['augmentations'][instr][el]360 361 # Channel shuffle362 if 'channel_shuffle' in augs:363 if augs['channel_shuffle'] > 0:364 if random.uniform(0, 1) < augs['channel_shuffle']:365 source = source[::-1].copy()366 applied_augs.append('channel_shuffle')367 # Random inverse368 if 'random_inverse' in augs:369 if augs['random_inverse'] > 0:370 if random.uniform(0, 1) < augs['random_inverse']:371 source = source[:, ::-1].copy()372 applied_augs.append('random_inverse')373 # Random polarity (multiply -1)374 if 'random_polarity' in augs:375 if augs['random_polarity'] > 0:376 if random.uniform(0, 1) < augs['random_polarity']:377 source = -source.copy()378 applied_augs.append('random_polarity')379 # Random pitch shift380 if 'pitch_shift' in augs:381 if augs['pitch_shift'] > 0:382 if random.uniform(0, 1) < augs['pitch_shift']:383 apply_aug = AU.PitchShift(384 min_semitones=augs['pitch_shift_min_semitones'],385 max_semitones=augs['pitch_shift_max_semitones'],386 p=1.0387 )388 source = apply_aug(samples=source, sample_rate=44100)389 applied_augs.append('pitch_shift')390 # Random seven band parametric eq391 if 'seven_band_parametric_eq' in augs:392 if augs['seven_band_parametric_eq'] > 0:393 if random.uniform(0, 1) < augs['seven_band_parametric_eq']:394 apply_aug = AU.SevenBandParametricEQ(395 min_gain_db=augs['seven_band_parametric_eq_min_gain_db'],396 max_gain_db=augs['seven_band_parametric_eq_max_gain_db'],397 p=1.0398 )399 source = apply_aug(samples=source, sample_rate=44100)400 applied_augs.append('seven_band_parametric_eq')401 # Random tanh distortion402 if 'tanh_distortion' in augs:403 if augs['tanh_distortion'] > 0:404 if random.uniform(0, 1) < augs['tanh_distortion']:405 apply_aug = AU.TanhDistortion(406 min_distortion=augs['tanh_distortion_min'],407 max_distortion=augs['tanh_distortion_max'],408 p=1.0409 )410 source = apply_aug(samples=source, sample_rate=44100)411 applied_augs.append('tanh_distortion')412 # Random MP3 Compression413 if 'mp3_compression' in augs:414 if augs['mp3_compression'] > 0:415 if random.uniform(0, 1) < augs['mp3_compression']:416 apply_aug = AU.Mp3Compression(417 min_bitrate=augs['mp3_compression_min_bitrate'],418 max_bitrate=augs['mp3_compression_max_bitrate'],419 backend=augs['mp3_compression_backend'],420 p=1.0421 )422 source = apply_aug(samples=source, sample_rate=44100)423 applied_augs.append('mp3_compression')424 # Random AddGaussianNoise425 if 'gaussian_noise' in augs:426 if augs['gaussian_noise'] > 0:427 if random.uniform(0, 1) < augs['gaussian_noise']:428 apply_aug = AU.AddGaussianNoise(429 min_amplitude=augs['gaussian_noise_min_amplitude'],430 max_amplitude=augs['gaussian_noise_max_amplitude'],431 p=1.0432 )433 source = apply_aug(samples=source, sample_rate=44100)434 applied_augs.append('gaussian_noise')435 # Random TimeStretch436 if 'time_stretch' in augs:437 if augs['time_stretch'] > 0:438 if random.uniform(0, 1) < augs['time_stretch']:439 apply_aug = AU.TimeStretch(440 min_rate=augs['time_stretch_min_rate'],441 max_rate=augs['time_stretch_max_rate'],442 leave_length_unchanged=True,443 p=1.0444 )445 source = apply_aug(samples=source, sample_rate=44100)446 applied_augs.append('time_stretch')447 448 # Possible fix of shape449 if source_shape != source.shape:450 source = source[..., :source_shape[-1]]451 452 # Random Reverb453 if 'pedalboard_reverb' in augs:454 if augs['pedalboard_reverb'] > 0:455 if random.uniform(0, 1) < augs['pedalboard_reverb']:456 room_size = random.uniform(457 augs['pedalboard_reverb_room_size_min'],458 augs['pedalboard_reverb_room_size_max'],459 )460 damping = random.uniform(461 augs['pedalboard_reverb_damping_min'],462 augs['pedalboard_reverb_damping_max'],463 )464 wet_level = random.uniform(465 augs['pedalboard_reverb_wet_level_min'],466 augs['pedalboard_reverb_wet_level_max'],467 )468 dry_level = random.uniform(469 augs['pedalboard_reverb_dry_level_min'],470 augs['pedalboard_reverb_dry_level_max'],471 )472 width = random.uniform(473 augs['pedalboard_reverb_width_min'],474 augs['pedalboard_reverb_width_max'],475 )476 board = PB.Pedalboard([PB.Reverb(477 room_size=room_size, # 0.1 - 0.9478 damping=damping, # 0.1 - 0.9479 wet_level=wet_level, # 0.1 - 0.9480 dry_level=dry_level, # 0.1 - 0.9481 width=width, # 0.9 - 1.0482 freeze_mode=0.0,483 )])484 source = board(source, 44100)485 applied_augs.append('pedalboard_reverb')486 487 # Random Chorus488 if 'pedalboard_chorus' in augs:489 if augs['pedalboard_chorus'] > 0:490 if random.uniform(0, 1) < augs['pedalboard_chorus']:491 rate_hz = random.uniform(492 augs['pedalboard_chorus_rate_hz_min'],493 augs['pedalboard_chorus_rate_hz_max'],494 )495 depth = random.uniform(496 augs['pedalboard_chorus_depth_min'],497 augs['pedalboard_chorus_depth_max'],498 )499 centre_delay_ms = random.uniform(500 augs['pedalboard_chorus_centre_delay_ms_min'],501 augs['pedalboard_chorus_centre_delay_ms_max'],502 )503 feedback = random.uniform(504 augs['pedalboard_chorus_feedback_min'],505 augs['pedalboard_chorus_feedback_max'],506 )507 mix = random.uniform(508 augs['pedalboard_chorus_mix_min'],509 augs['pedalboard_chorus_mix_max'],510 )511 board = PB.Pedalboard([PB.Chorus(512 rate_hz=rate_hz,513 depth=depth,514 centre_delay_ms=centre_delay_ms,515 feedback=feedback,516 mix=mix,517 )])518 source = board(source, 44100)519 applied_augs.append('pedalboard_chorus')520 521 # Random Phazer522 if 'pedalboard_phazer' in augs:523 if augs['pedalboard_phazer'] > 0:524 if random.uniform(0, 1) < augs['pedalboard_phazer']:525 rate_hz = random.uniform(526 augs['pedalboard_phazer_rate_hz_min'],527 augs['pedalboard_phazer_rate_hz_max'],528 )529 depth = random.uniform(530 augs['pedalboard_phazer_depth_min'],531 augs['pedalboard_phazer_depth_max'],532 )533 centre_frequency_hz = random.uniform(534 augs['pedalboard_phazer_centre_frequency_hz_min'],535 augs['pedalboard_phazer_centre_frequency_hz_max'],536 )537 feedback = random.uniform(538 augs['pedalboard_phazer_feedback_min'],539 augs['pedalboard_phazer_feedback_max'],540 )541 mix = random.uniform(542 augs['pedalboard_phazer_mix_min'],543 augs['pedalboard_phazer_mix_max'],544 )545 board = PB.Pedalboard([PB.Phaser(546 rate_hz=rate_hz,547 depth=depth,548 centre_frequency_hz=centre_frequency_hz,549 feedback=feedback,550 mix=mix,551 )])552 source = board(source, 44100)553 applied_augs.append('pedalboard_phazer')554 555 # Random Distortion556 if 'pedalboard_distortion' in augs:557 if augs['pedalboard_distortion'] > 0:558 if random.uniform(0, 1) < augs['pedalboard_distortion']:559 drive_db = random.uniform(560 augs['pedalboard_distortion_drive_db_min'],561 augs['pedalboard_distortion_drive_db_max'],562 )563 board = PB.Pedalboard([PB.Distortion(564 drive_db=drive_db,565 )])566 source = board(source, 44100)567 applied_augs.append('pedalboard_distortion')568 569 # Random PitchShift570 if 'pedalboard_pitch_shift' in augs:571 if augs['pedalboard_pitch_shift'] > 0:572 if random.uniform(0, 1) < augs['pedalboard_pitch_shift']:573 semitones = random.uniform(574 augs['pedalboard_pitch_shift_semitones_min'],575 augs['pedalboard_pitch_shift_semitones_max'],576 )577 board = PB.Pedalboard([PB.PitchShift(578 semitones=semitones579 )])580 source = board(source, 44100)581 applied_augs.append('pedalboard_pitch_shift')582 583 # Random Resample584 if 'pedalboard_resample' in augs:585 if augs['pedalboard_resample'] > 0:586 if random.uniform(0, 1) < augs['pedalboard_resample']:587 target_sample_rate = random.uniform(588 augs['pedalboard_resample_target_sample_rate_min'],589 augs['pedalboard_resample_target_sample_rate_max'],590 )591 board = PB.Pedalboard([PB.Resample(592 target_sample_rate=target_sample_rate593 )])594 source = board(source, 44100)595 applied_augs.append('pedalboard_resample')596 597 # Random Bitcrash598 if 'pedalboard_bitcrash' in augs:599 if augs['pedalboard_bitcrash'] > 0:600 if random.uniform(0, 1) < augs['pedalboard_bitcrash']:601 bit_depth = random.uniform(602 augs['pedalboard_bitcrash_bit_depth_min'],603 augs['pedalboard_bitcrash_bit_depth_max'],604 )605 board = PB.Pedalboard([PB.Bitcrush(606 bit_depth=bit_depth607 )])608 source = board(source, 44100)609 applied_augs.append('pedalboard_bitcrash')610 611 # Random MP3Compressor612 if 'pedalboard_mp3_compressor' in augs:613 if augs['pedalboard_mp3_compressor'] > 0:614 if random.uniform(0, 1) < augs['pedalboard_mp3_compressor']:615 vbr_quality = random.uniform(616 augs['pedalboard_mp3_compressor_pedalboard_mp3_compressor_min'],617 augs['pedalboard_mp3_compressor_pedalboard_mp3_compressor_max'],618 )619 board = PB.Pedalboard([PB.MP3Compressor(620 vbr_quality=vbr_quality621 )])622 source = board(source, 44100)623 applied_augs.append('pedalboard_mp3_compressor')624 625 # print(applied_augs)626 return source627 628 def __getitem__(self, index):629 if self.dataset_type in [1, 2, 3]:630 res = self.load_random_mix()631 else:632 res = self.load_aligned_data()633 634 # Randomly change loudness of each stem635 if self.aug:636 if 'loudness' in self.config['augmentations']:637 if self.config['augmentations']['loudness']:638 loud_values = np.random.uniform(639 low=self.config['augmentations']['loudness_min'],640 high=self.config['augmentations']['loudness_max'],641 size=(len(res),)642 )643 loud_values = torch.tensor(loud_values, dtype=torch.float32)644 res *= loud_values[:, None, None]645 646 mix = res.sum(0)647 648 if self.aug:649 if 'mp3_compression_on_mixture' in self.config['augmentations']:650 apply_aug = AU.Mp3Compression(651 min_bitrate=self.config['augmentations']['mp3_compression_on_mixture_bitrate_min'],652 max_bitrate=self.config['augmentations']['mp3_compression_on_mixture_bitrate_max'],653 backend=self.config['augmentations']['mp3_compression_on_mixture_backend'],654 p=self.config['augmentations']['mp3_compression_on_mixture']655 )656 mix_conv = mix.cpu().numpy().astype(np.float32)657 required_shape = mix_conv.shape658 mix = apply_aug(samples=mix_conv, sample_rate=44100)659 # Sometimes it gives longer audio (so we cut)660 if mix.shape != required_shape:661 mix = mix[..., :required_shape[-1]]662 mix = torch.tensor(mix, dtype=torch.float32)663 664 # If we need to optimize only given stem665 if self.config.training.target_instrument is not None:666 index = self.config.training.instruments.index(self.config.training.target_instrument)667 return res[index:index+1], mix668 669 return res, mix670