Spaces:
Running
Running
| # Adapted from: | |
| # https://github.com/csteinmetz1/micro-tcn/blob/main/microtcn/utils.py | |
| import os | |
| import csv | |
| import torch | |
| import fnmatch | |
| import numpy as np | |
| import random | |
| from enum import Enum | |
| import pyloudnorm as pyln | |
| class DSPMode(Enum): | |
| NONE = "none" | |
| TRAIN_INFER = "train_infer" | |
| INFER = "infer" | |
| def __str__(self): | |
| return self.value | |
| def loudness_normalize(x, sample_rate, target_loudness=-24.0): | |
| x = x.view(1, -1) | |
| stereo_audio = x.repeat(2, 1).permute(1, 0).numpy() | |
| meter = pyln.Meter(sample_rate) | |
| loudness = meter.integrated_loudness(stereo_audio) | |
| norm_x = pyln.normalize.loudness( | |
| stereo_audio, | |
| loudness, | |
| target_loudness, | |
| ) | |
| x = torch.tensor(norm_x).permute(1, 0) | |
| x = x[0, :].view(1, -1) | |
| return x | |
| def get_random_file_id(keys): | |
| # generate a random index into the keys of the input files | |
| rand_input_idx = torch.randint(0, len(keys) - 1, [1])[0] | |
| # find the key (file_id) correponding to the random index | |
| rand_input_file_id = list(keys)[rand_input_idx] | |
| return rand_input_file_id | |
| def get_random_patch(audio_file, length, check_silence=True): | |
| silent = True | |
| while silent: | |
| start_idx = int(torch.rand(1) * (audio_file.num_frames - length)) | |
| stop_idx = start_idx + length | |
| patch = audio_file.audio[:, start_idx:stop_idx].clone().detach() | |
| if (patch ** 2).mean() > 1e-4 or not check_silence: | |
| silent = False | |
| return start_idx, stop_idx | |
| def seed_worker(worker_id): | |
| worker_seed = torch.initial_seed() % 2 ** 32 | |
| np.random.seed(worker_seed) | |
| random.seed(worker_seed) | |
| def getFilesPath(directory, extension): | |
| n_path = [] | |
| for path, subdirs, files in os.walk(directory): | |
| for name in files: | |
| if fnmatch.fnmatch(name, extension): | |
| n_path.append(os.path.join(path, name)) | |
| n_path.sort() | |
| return n_path | |
| def count_parameters(model, trainable_only=True): | |
| if trainable_only: | |
| if len(list(model.parameters())) > 0: | |
| params = sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| else: | |
| params = 0 | |
| else: | |
| if len(list(model.parameters())) > 0: | |
| params = sum(p.numel() for p in model.parameters()) | |
| else: | |
| params = 0 | |
| return params | |
| def system_summary(system): | |
| print(f"Encoder: {count_parameters(system.encoder)/1e6:0.2f} M") | |
| print(f"Processor: {count_parameters(system.processor)/1e6:0.2f} M") | |
| if hasattr(system, "adv_loss_fn"): | |
| for idx, disc in enumerate(system.adv_loss_fn.discriminators): | |
| print(f"Discriminator {idx+1}: {count_parameters(disc)/1e6:0.2f} M") | |
| def center_crop(x, length: int): | |
| if x.shape[-1] != length: | |
| start = (x.shape[-1] - length) // 2 | |
| stop = start + length | |
| x = x[..., start:stop] | |
| return x | |
| def causal_crop(x, length: int): | |
| if x.shape[-1] != length: | |
| stop = x.shape[-1] - 1 | |
| start = stop - length | |
| x = x[..., start:stop] | |
| return x | |
| def denormalize(norm_val, max_val, min_val): | |
| return (norm_val * (max_val - min_val)) + min_val | |
| def normalize(denorm_val, max_val, min_val): | |
| return (denorm_val - min_val) / (max_val - min_val) | |
| def get_random_patch(audio_file, length, energy_treshold=1e-4): | |
| """Produce sample indicies for a random patch of size `length`. | |
| This function will check the energy of the selected patch to | |
| ensure that it is not complete silence. If silence is found, | |
| it will continue searching for a non-silent patch. | |
| Args: | |
| audio_file (AudioFile): Audio file object. | |
| length (int): Number of samples in random patch. | |
| Returns: | |
| start_idx (int): Starting sample index | |
| stop_idx (int): Stop sample index | |
| """ | |
| silent = True | |
| while silent: | |
| start_idx = int(torch.rand(1) * (audio_file.num_frames - length)) | |
| stop_idx = start_idx + length | |
| patch = audio_file.audio[:, start_idx:stop_idx] | |
| if (patch ** 2).mean() > energy_treshold: | |
| silent = False | |
| return start_idx, stop_idx | |
| def split_dataset(file_list, subset, train_frac): | |
| """Given a list of files, split into train/val/test sets. | |
| Args: | |
| file_list (list): List of audio files. | |
| subset (str): One of "train", "val", or "test". | |
| train_frac (float): Fraction of the dataset to use for training. | |
| Returns: | |
| file_list (list): List of audio files corresponding to subset. | |
| """ | |
| assert train_frac > 0.1 and train_frac < 1.0 | |
| total_num_examples = len(file_list) | |
| train_num_examples = int(total_num_examples * train_frac) | |
| val_num_examples = int(total_num_examples * (1 - train_frac) / 2) | |
| test_num_examples = total_num_examples - (train_num_examples + val_num_examples) | |
| if train_num_examples < 0: | |
| raise ValueError( | |
| f"No examples in training set. Try increasing train_frac: {train_frac}." | |
| ) | |
| elif val_num_examples < 0: | |
| raise ValueError( | |
| f"No examples in validation set. Try decreasing train_frac: {train_frac}." | |
| ) | |
| elif test_num_examples < 0: | |
| raise ValueError( | |
| f"No examples in test set. Try decreasing train_frac: {train_frac}." | |
| ) | |
| if subset == "train": | |
| start_idx = 0 | |
| stop_idx = train_num_examples | |
| elif subset == "val": | |
| start_idx = train_num_examples | |
| stop_idx = start_idx + val_num_examples | |
| elif subset == "test": | |
| start_idx = train_num_examples + val_num_examples | |
| stop_idx = start_idx + test_num_examples + 1 | |
| else: | |
| raise ValueError("Invalid subset: {subset}.") | |
| return file_list[start_idx:stop_idx] | |
| def rademacher(size): | |
| """Generates random samples from a Rademacher distribution +-1 | |
| Args: | |
| size (int): | |
| """ | |
| m = torch.distributions.binomial.Binomial(1, 0.5) | |
| x = m.sample(size) | |
| x[x == 0] = -1 | |
| return x | |
| def get_subset(csv_file): | |
| subset_files = [] | |
| with open(csv_file) as fp: | |
| reader = csv.DictReader(fp) | |
| for row in reader: | |
| subset_files.append(row["filepath"]) | |
| return list(set(subset_files)) | |
| def conform_length(x: torch.Tensor, length: int): | |
| """Crop or pad input on last dim to match `length`.""" | |
| if x.shape[-1] < length: | |
| padsize = length - x.shape[-1] | |
| x = torch.nn.functional.pad(x, (0, padsize)) | |
| elif x.shape[-1] > length: | |
| x = x[..., :length] | |
| return x | |
| def linear_fade( | |
| x: torch.Tensor, | |
| fade_ms: float = 50.0, | |
| sample_rate: float = 22050, | |
| ): | |
| """Apply fade in and fade out to last dim.""" | |
| fade_samples = int(fade_ms * 1e-3 * 22050) | |
| fade_in = torch.linspace(0.0, 1.0, steps=fade_samples) | |
| fade_out = torch.linspace(1.0, 0.0, steps=fade_samples) | |
| # fade in | |
| x[..., :fade_samples] *= fade_in | |
| # fade out | |
| x[..., -fade_samples:] *= fade_out | |
| return x | |
| # def get_random_patch(x, sample_rate, length_samples): | |
| # length = length_samples | |
| # silent = True | |
| # while silent: | |
| # start_idx = np.random.randint(0, x.shape[-1] - length - 1) | |
| # stop_idx = start_idx + length | |
| # x_crop = x[0:1, start_idx:stop_idx] | |
| # # check for silence | |
| # frames = length // sample_rate | |
| # silent_frames = [] | |
| # for n in range(frames): | |
| # start_idx = n * sample_rate | |
| # stop_idx = start_idx + sample_rate | |
| # x_frame = x_crop[0:1, start_idx:stop_idx] | |
| # if (x_frame ** 2).mean() > 3e-4: | |
| # silent_frames.append(False) | |
| # else: | |
| # silent_frames.append(True) | |
| # silent = True if any(silent_frames) else False | |
| # x_crop /= x_crop.abs().max() | |
| # return x_crop | |