| |
| __author__ = "Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/" |
|
|
| import os |
| import librosa |
| import soundfile as sf |
| import numpy as np |
| import argparse |
|
|
|
|
| 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) |
| |
| 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) |
|
|
