Spaces:
Runtime error
Runtime error
| # coding: utf-8 | |
| __author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" | |
| import argparse | |
| import os | |
| import librosa | |
| import numpy as np | |
| import soundfile as sf | |
| def stft(wave, nfft, hl): | |
| wave_left = np.asfortranarray(wave[0]) | |
| wave_right = np.asfortranarray(wave[1]) | |
| spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl) | |
| spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl) | |
| spec = np.asfortranarray([spec_left, spec_right]) | |
| return spec | |
| def istft(spec, hl, length): | |
| spec_left = np.asfortranarray(spec[0]) | |
| spec_right = np.asfortranarray(spec[1]) | |
| wave_left = librosa.istft(spec_left, hop_length=hl, length=length) | |
| wave_right = librosa.istft(spec_right, hop_length=hl, length=length) | |
| wave = np.asfortranarray([wave_left, wave_right]) | |
| return wave | |
| def absmax(a, *, axis): | |
| dims = list(a.shape) | |
| dims.pop(axis) | |
| indices = np.ogrid[tuple(slice(0, d) for d in dims)] | |
| argmax = np.abs(a).argmax(axis=axis) | |
| indices.insert((len(a.shape) + axis) % len(a.shape), argmax) | |
| return a[tuple(indices)] | |
| def absmin(a, *, axis): | |
| dims = list(a.shape) | |
| dims.pop(axis) | |
| indices = np.ogrid[tuple(slice(0, d) for d in dims)] | |
| argmax = np.abs(a).argmin(axis=axis) | |
| indices.insert((len(a.shape) + axis) % len(a.shape), argmax) | |
| return a[tuple(indices)] | |
| def lambda_max(arr, axis=None, key=None, keepdims=False): | |
| idxs = np.argmax(key(arr), axis) | |
| if axis is not None: | |
| idxs = np.expand_dims(idxs, axis) | |
| result = np.take_along_axis(arr, idxs, axis) | |
| if not keepdims: | |
| result = np.squeeze(result, axis=axis) | |
| return result | |
| else: | |
| return arr.flatten()[idxs] | |
| def lambda_min(arr, axis=None, key=None, keepdims=False): | |
| idxs = np.argmin(key(arr), axis) | |
| if axis is not None: | |
| idxs = np.expand_dims(idxs, axis) | |
| result = np.take_along_axis(arr, idxs, axis) | |
| if not keepdims: | |
| result = np.squeeze(result, axis=axis) | |
| return result | |
| else: | |
| return arr.flatten()[idxs] | |
| def average_waveforms(pred_track, weights, algorithm): | |
| """ | |
| :param pred_track: shape = (num, channels, length) | |
| :param weights: shape = (num, ) | |
| :param algorithm: One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft | |
| :return: averaged waveform in shape (channels, length) | |
| """ | |
| pred_track = np.array(pred_track) | |
| final_length = pred_track.shape[-1] | |
| mod_track = [] | |
| for i in range(pred_track.shape[0]): | |
| if algorithm == "avg_wave": | |
| mod_track.append(pred_track[i] * weights[i]) | |
| elif algorithm in ["median_wave", "min_wave", "max_wave"]: | |
| mod_track.append(pred_track[i]) | |
| elif algorithm in ["avg_fft", "min_fft", "max_fft", "median_fft"]: | |
| spec = stft(pred_track[i], nfft=2048, hl=1024) | |
| if algorithm in ["avg_fft"]: | |
| mod_track.append(spec * weights[i]) | |
| else: | |
| mod_track.append(spec) | |
| pred_track = np.array(mod_track) | |
| if algorithm in ["avg_wave"]: | |
| pred_track = pred_track.sum(axis=0) | |
| pred_track /= np.array(weights).sum().T | |
| elif algorithm in ["median_wave"]: | |
| pred_track = np.median(pred_track, axis=0) | |
| elif algorithm in ["min_wave"]: | |
| pred_track = np.array(pred_track) | |
| pred_track = lambda_min(pred_track, axis=0, key=np.abs) | |
| elif algorithm in ["max_wave"]: | |
| pred_track = np.array(pred_track) | |
| pred_track = lambda_max(pred_track, axis=0, key=np.abs) | |
| elif algorithm in ["avg_fft"]: | |
| pred_track = pred_track.sum(axis=0) | |
| pred_track /= np.array(weights).sum() | |
| pred_track = istft(pred_track, 1024, final_length) | |
| elif algorithm in ["min_fft"]: | |
| pred_track = np.array(pred_track) | |
| pred_track = lambda_min(pred_track, axis=0, key=np.abs) | |
| pred_track = istft(pred_track, 1024, final_length) | |
| elif algorithm in ["max_fft"]: | |
| pred_track = np.array(pred_track) | |
| pred_track = absmax(pred_track, axis=0) | |
| pred_track = istft(pred_track, 1024, final_length) | |
| elif algorithm in ["median_fft"]: | |
| pred_track = np.array(pred_track) | |
| pred_track = np.median(pred_track, axis=0) | |
| pred_track = istft(pred_track, 1024, final_length) | |
| return pred_track | |
| def ensemble_files(args): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--files", | |
| type=str, | |
| required=True, | |
| nargs="+", | |
| help="Path to all audio-files to ensemble", | |
| ) | |
| parser.add_argument( | |
| "--type", | |
| type=str, | |
| default="avg_wave", | |
| help="One of avg_wave, median_wave, min_wave, max_wave, avg_fft, median_fft, min_fft, max_fft", | |
| ) | |
| parser.add_argument( | |
| "--weights", | |
| type=float, | |
| nargs="+", | |
| help="Weights to create ensemble. Number of weights must be equal to number of files", | |
| ) | |
| parser.add_argument( | |
| "--output", | |
| default="res.wav", | |
| type=str, | |
| help="Path to wav file where ensemble result will be stored", | |
| ) | |
| if args is None: | |
| args = parser.parse_args() | |
| else: | |
| args = parser.parse_args(args) | |
| print("Ensemble type: {}".format(args.type)) | |
| print("Number of input files: {}".format(len(args.files))) | |
| if args.weights is not None: | |
| weights = args.weights | |
| else: | |
| weights = np.ones(len(args.files)) | |
| print("Weights: {}".format(weights)) | |
| print("Output file: {}".format(args.output)) | |
| data = [] | |
| for f in args.files: | |
| if not os.path.isfile(f): | |
| print("Error. Can't find file: {}. Check paths.".format(f)) | |
| exit() | |
| print("Reading file: {}".format(f)) | |
| wav, sr = librosa.load(f, sr=None, mono=False) | |
| # wav, sr = sf.read(f) | |
| print("Waveform shape: {} sample rate: {}".format(wav.shape, sr)) | |
| data.append(wav) | |
| data = np.array(data) | |
| res = average_waveforms(data, weights, args.type) | |
| print("Result shape: {}".format(res.shape)) | |
| sf.write(args.output, res.T, sr, "FLOAT") | |
| if __name__ == "__main__": | |
| ensemble_files(None) | |