ultimatevocalremovergui/lib/spec_utils.py
2020-07-20 16:54:03 -05:00

137 lines
4.5 KiB
Python

import os
import librosa
import numpy as np
import soundfile as sf
import torch
def crop_center(h1, h2, concat=True):
# s_freq = (h2.shape[2] - h1.shape[2]) // 2
# e_freq = s_freq + h1.shape[2]
h1_shape = h1.size()
h2_shape = h2.size()
if h2_shape[3] < h1_shape[3]:
raise ValueError('h2_shape[3] must be greater than h1_shape[3]')
s_time = (h2_shape[3] - h1_shape[3]) // 2
e_time = s_time + h1_shape[3]
h2 = h2[:, :, :, s_time:e_time]
if concat:
return torch.cat([h1, h2], dim=1)
else:
return h2
def calc_spec(X, hop_length):
n_fft = (hop_length - 1) * 2
audio_left = np.asfortranarray(X[0])
audio_right = np.asfortranarray(X[1])
spec_left = librosa.stft(audio_left, n_fft, hop_length=hop_length)
spec_right = librosa.stft(audio_right, n_fft, hop_length=hop_length)
spec = np.asfortranarray([spec_left, spec_right])
return spec
def mask_uninformative(mask, ref, thres=0.3, min_range=64, fade_area=32):
if min_range < fade_area * 2:
raise ValueError('min_range must be >= fade_area * 2')
idx = np.where(ref.mean(axis=(0, 1)) < thres)[0]
starts = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0])
ends = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1])
uninformative = np.where(ends - starts > min_range)[0]
if len(uninformative) > 0:
starts = starts[uninformative]
ends = ends[uninformative]
old_e = None
for s, e in zip(starts, ends):
if old_e is not None and s - old_e < fade_area:
s = old_e - fade_area * 2
elif s != 0:
start_mask = mask[:, :, s:s + fade_area]
np.clip(
start_mask + np.linspace(0, 1, fade_area), 0, 1,
out=start_mask)
if e != mask.shape[2]:
end_mask = mask[:, :, e - fade_area:e]
np.clip(
end_mask + np.linspace(1, 0, fade_area), 0, 1,
out=end_mask)
mask[:, :, s + fade_area:e - fade_area] = 1
old_e = e
return mask
def align_wave_head_and_tail(a, b, sr):
a_mono = a[:, :sr * 4].sum(axis=0)
b_mono = b[:, :sr * 4].sum(axis=0)
a_mono -= a_mono.mean()
b_mono -= b_mono.mean()
offset = len(a_mono) - 1
delay = np.argmax(np.correlate(a_mono, b_mono, 'full')) - offset
if delay > 0:
a = a[:, delay:]
else:
b = b[:, np.abs(delay):]
if a.shape[1] < b.shape[1]:
b = b[:, :a.shape[1]]
else:
a = a[:, :b.shape[1]]
return a, b
def cache_or_load(mix_path, inst_path, sr, hop_length):
_, mix_ext = os.path.splitext(mix_path)
_, inst_ext = os.path.splitext(inst_path)
spec_mix_path = mix_path.replace(mix_ext, '.npy')
spec_inst_path = inst_path.replace(inst_ext, '.npy')
if os.path.exists(spec_mix_path) and os.path.exists(spec_inst_path):
X = np.load(spec_mix_path)
y = np.load(spec_inst_path)
else:
X, _ = librosa.load(
mix_path, sr, False, dtype=np.float32, res_type='kaiser_fast')
y, _ = librosa.load(
inst_path, sr, False, dtype=np.float32, res_type='kaiser_fast')
X, _ = librosa.effects.trim(X)
y, _ = librosa.effects.trim(y)
X, y = align_wave_head_and_tail(X, y, sr)
X = np.abs(calc_spec(X, hop_length))
y = np.abs(calc_spec(y, hop_length))
_, ext = os.path.splitext(mix_path)
np.save(spec_mix_path, X)
np.save(spec_inst_path, y)
return X, y
def spec_to_wav(mag, phase, hop_length):
spec = mag * phase
spec_left = np.asfortranarray(spec[0])
spec_right = np.asfortranarray(spec[1])
wav_left = librosa.istft(spec_left, hop_length=hop_length)
wav_right = librosa.istft(spec_right, hop_length=hop_length)
wav = np.asfortranarray([wav_left, wav_right])
return wav
if __name__ == "__main__":
import sys
X, _ = librosa.load(
sys.argv[1], 44100, False, dtype=np.float32, res_type='kaiser_fast')
y, _ = librosa.load(
sys.argv[2], 44100, False, dtype=np.float32, res_type='kaiser_fast')
X, _ = librosa.effects.trim(X)
y, _ = librosa.effects.trim(y)
X, y = align_wave_head_and_tail(X, y, 44100)
sf.write('test_i.wav', y.T, 44100)
sf.write('test_m.wav', X.T, 44100)
sf.write('test_v.wav', (X - y).T, 44100)