Spaces:
Runtime error
Runtime error
| import argparse | |
| import os | |
| import librosa | |
| import numpy as np | |
| import soundfile as sf | |
| import torch | |
| from tqdm import tqdm | |
| from lib import dataset | |
| from lib import nets | |
| from lib import spec_utils | |
| from lib import utils | |
| class Separator(object): | |
| def __init__(self, model, device, batchsize, cropsize, postprocess=False): | |
| self.model = model | |
| self.offset = model.offset | |
| self.device = device | |
| self.batchsize = batchsize | |
| self.cropsize = cropsize | |
| self.postprocess = postprocess | |
| def _separate(self, X_mag_pad, roi_size): | |
| X_dataset = [] | |
| patches = (X_mag_pad.shape[2] - 2 * self.offset) // roi_size | |
| for i in range(patches): | |
| start = i * roi_size | |
| X_mag_crop = X_mag_pad[:, :, start:start + self.cropsize] | |
| X_dataset.append(X_mag_crop) | |
| X_dataset = np.asarray(X_dataset) | |
| self.model.eval() | |
| with torch.no_grad(): | |
| mask = [] | |
| # To reduce the overhead, dataloader is not used. | |
| for i in tqdm(range(0, patches, self.batchsize)): | |
| X_batch = X_dataset[i: i + self.batchsize] | |
| X_batch = torch.from_numpy(X_batch).to(self.device) | |
| pred = self.model.predict_mask(X_batch) | |
| pred = pred.detach().cpu().numpy() | |
| pred = np.concatenate(pred, axis=2) | |
| mask.append(pred) | |
| mask = np.concatenate(mask, axis=2) | |
| return mask | |
| def _preprocess(self, X_spec): | |
| X_mag = np.abs(X_spec) | |
| X_phase = np.angle(X_spec) | |
| return X_mag, X_phase | |
| def _postprocess(self, mask, X_mag, X_phase): | |
| if self.postprocess: | |
| mask = spec_utils.merge_artifacts(mask) | |
| y_spec = mask * X_mag * np.exp(1.j * X_phase) | |
| v_spec = (1 - mask) * X_mag * np.exp(1.j * X_phase) | |
| return y_spec, v_spec | |
| def separate(self, X_spec): | |
| X_mag, X_phase = self._preprocess(X_spec) | |
| n_frame = X_mag.shape[2] | |
| pad_l, pad_r, roi_size = dataset.make_padding(n_frame, self.cropsize, self.offset) | |
| X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant') | |
| X_mag_pad /= X_mag_pad.max() | |
| mask = self._separate(X_mag_pad, roi_size) | |
| mask = mask[:, :, :n_frame] | |
| y_spec, v_spec = self._postprocess(mask, X_mag, X_phase) | |
| return y_spec, v_spec | |
| def separate_tta(self, X_spec): | |
| X_mag, X_phase = self._preprocess(X_spec) | |
| n_frame = X_mag.shape[2] | |
| pad_l, pad_r, roi_size = dataset.make_padding(n_frame, self.cropsize, self.offset) | |
| X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant') | |
| X_mag_pad /= X_mag_pad.max() | |
| mask = self._separate(X_mag_pad, roi_size) | |
| pad_l += roi_size // 2 | |
| pad_r += roi_size // 2 | |
| X_mag_pad = np.pad(X_mag, ((0, 0), (0, 0), (pad_l, pad_r)), mode='constant') | |
| X_mag_pad /= X_mag_pad.max() | |
| mask_tta = self._separate(X_mag_pad, roi_size) | |
| mask_tta = mask_tta[:, :, roi_size // 2:] | |
| mask = (mask[:, :, :n_frame] + mask_tta[:, :, :n_frame]) * 0.5 | |
| y_spec, v_spec = self._postprocess(mask, X_mag, X_phase) | |
| return y_spec, v_spec | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument('--gpu', '-g', type=int, default=-1) | |
| p.add_argument('--pretrained_model', '-P', type=str, default='models/baseline.pth') | |
| p.add_argument('--input', '-i', required=True) | |
| p.add_argument('--sr', '-r', type=int, default=44100) | |
| p.add_argument('--n_fft', '-f', type=int, default=2048) | |
| p.add_argument('--hop_length', '-H', type=int, default=1024) | |
| p.add_argument('--batchsize', '-B', type=int, default=4) | |
| p.add_argument('--cropsize', '-c', type=int, default=256) | |
| p.add_argument('--output_image', '-I', action='store_true') | |
| p.add_argument('--postprocess', '-p', action='store_true') | |
| p.add_argument('--tta', '-t', action='store_true') | |
| p.add_argument('--output_dir', '-o', type=str, default="") | |
| args = p.parse_args() | |
| print('loading model...', end=' ') | |
| device = torch.device('cpu') | |
| model = nets.CascadedNet(args.n_fft, 32, 128) | |
| model.load_state_dict(torch.load(args.pretrained_model, map_location=device)) | |
| if torch.cuda.is_available() and args.gpu >= 0: | |
| device = torch.device('cuda:{}'.format(args.gpu)) | |
| model.to(device) | |
| print('done') | |
| print('loading wave source...', end=' ') | |
| X, sr = librosa.load( | |
| args.input, args.sr, False, dtype=np.float32, res_type='kaiser_fast') | |
| basename = os.path.splitext(os.path.basename(args.input))[0] | |
| print('done') | |
| if X.ndim == 1: | |
| # mono to stereo | |
| X = np.asarray([X, X]) | |
| print('stft of wave source...', end=' ') | |
| X_spec = spec_utils.wave_to_spectrogram(X, args.hop_length, args.n_fft) | |
| print('done') | |
| sp = Separator(model, device, args.batchsize, args.cropsize, args.postprocess) | |
| if args.tta: | |
| y_spec, v_spec = sp.separate_tta(X_spec) | |
| else: | |
| y_spec, v_spec = sp.separate(X_spec) | |
| print('validating output directory...', end=' ') | |
| output_dir = args.output_dir | |
| if output_dir != "": # modifies output_dir if theres an arg specified | |
| output_dir = output_dir.rstrip('/') + '/' | |
| os.makedirs(output_dir, exist_ok=True) | |
| print('done') | |
| print('inverse stft of instruments...', end=' ') | |
| wave = spec_utils.spectrogram_to_wave(y_spec, hop_length=args.hop_length) | |
| print('done') | |
| # sf.write('{}{}_Instruments.wav'.format(output_dir, basename), wave.T, sr) | |
| sf.write('{}Instruments.wav'.format(output_dir), wave.T, sr) | |
| print('inverse stft of vocals...', end=' ') | |
| wave = spec_utils.spectrogram_to_wave(v_spec, hop_length=args.hop_length) | |
| print('done') | |
| sf.write('{}{}_Vocals.wav'.format(output_dir, basename), wave.T, sr) | |
| if args.output_image: | |
| image = spec_utils.spectrogram_to_image(y_spec) | |
| utils.imwrite('{}{}_Instruments.jpg'.format(output_dir, basename), image) | |
| image = spec_utils.spectrogram_to_image(v_spec) | |
| utils.imwrite('{}{}_Vocals.jpg'.format(output_dir, basename), image) | |
| if __name__ == '__main__': | |
| main() | |