Duplicate from antonbol/vocal_remover

Co-authored-by: Anton Bölenius <antonbol@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Vocal Remover
+emoji: ⚡
+colorFrom: red
+colorTo: blue
+sdk: gradio
+sdk_version: 3.16.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: antonbol/vocal_remover
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,23 @@
+import gradio as gr
+import hopsworks
+import subprocess
+def vocal_remove(audio):    
+    project = hopsworks.login()
+    mr = project.get_model_registry()
+    # model = mr.get_best_model("vocal_remover", "validation_loss", "min")
+    model = mr.get_model("vocal_remover", version=3)
+    model_path = model.download()
+    model_path_pth = model_path + "/vocal_model.pth"
+    # print("model_path: ", model_path)s
+    subprocess.run(["python3", "inference.py", "--input", audio, "--pretrained_model", model_path_pth, "--output_dir", "./"])
+    return "./Instruments.wav"
+
+iface = gr.Interface(
+    fn=vocal_remove, 
+    inputs=gr.Audio(source="upload", type="filepath"), 
+    outputs="audio",
+    title="Vocal Remover",
+    description="Removes Vocals from song, currently undertrained, fragments of vocals can remain depending on song",
+)
+
+iface.launch()

inference.py

@@ -0,0 +1,181 @@
+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()

lib/dataset.py

@@ -0,0 +1,257 @@
+import os
+import random
+
+import numpy as np
+import torch
+import torch.utils.data
+from tqdm import tqdm
+
+try:
+    from lib import spec_utils
+except ModuleNotFoundError:
+    import spec_utils
+
+
+class VocalRemoverTrainingSet(torch.utils.data.Dataset):
+
+    def __init__(self, training_set, cropsize, reduction_rate, reduction_weight, mixup_rate, mixup_alpha):
+        self.training_set = training_set
+        self.cropsize = cropsize
+        self.reduction_rate = reduction_rate
+        self.reduction_weight = reduction_weight
+        self.mixup_rate = mixup_rate
+        self.mixup_alpha = mixup_alpha
+
+    def __len__(self):
+        return len(self.training_set)
+
+    def do_crop(self, X_path, y_path):
+        X_mmap = np.load(X_path, mmap_mode='r')
+        y_mmap = np.load(y_path, mmap_mode='r')
+
+        start = np.random.randint(0, X_mmap.shape[2] - self.cropsize)
+        end = start + self.cropsize
+
+        X_crop = np.array(X_mmap[:, :, start:end], copy=True)
+        y_crop = np.array(y_mmap[:, :, start:end], copy=True)
+
+        return X_crop, y_crop
+
+    def do_aug(self, X, y):
+        if np.random.uniform() < self.reduction_rate:
+            y = spec_utils.aggressively_remove_vocal(X, y, self.reduction_weight)
+
+        if np.random.uniform() < 0.5:
+            # swap channel
+            X = X[::-1].copy()
+            y = y[::-1].copy()
+
+        if np.random.uniform() < 0.01:
+            # inst
+            X = y.copy()
+
+        # if np.random.uniform() < 0.01:
+        #     # mono
+        #     X[:] = X.mean(axis=0, keepdims=True)
+        #     y[:] = y.mean(axis=0, keepdims=True)
+
+        return X, y
+
+    def do_mixup(self, X, y):
+        idx = np.random.randint(0, len(self))
+        X_path, y_path, coef = self.training_set[idx]
+
+        X_i, y_i = self.do_crop(X_path, y_path)
+        X_i /= coef
+        y_i /= coef
+
+        X_i, y_i = self.do_aug(X_i, y_i)
+
+        lam = np.random.beta(self.mixup_alpha, self.mixup_alpha)
+        X = lam * X + (1 - lam) * X_i
+        y = lam * y + (1 - lam) * y_i
+
+        return X, y
+
+    def __getitem__(self, idx):
+        X_path, y_path, coef = self.training_set[idx]
+
+        X, y = self.do_crop(X_path, y_path)
+        X /= coef
+        y /= coef
+
+        X, y = self.do_aug(X, y)
+
+        if np.random.uniform() < self.mixup_rate:
+            X, y = self.do_mixup(X, y)
+
+        X_mag = np.abs(X)
+        y_mag = np.abs(y)
+
+        return X_mag, y_mag
+
+
+class VocalRemoverValidationSet(torch.utils.data.Dataset):
+
+    def __init__(self, patch_list):
+        self.patch_list = patch_list
+
+    def __len__(self):
+        return len(self.patch_list)
+
+    def __getitem__(self, idx):
+        path = self.patch_list[idx]
+        data = np.load(path)
+
+        X, y = data['X'], data['y']
+
+        X_mag = np.abs(X)
+        y_mag = np.abs(y)
+
+        return X_mag, y_mag
+
+
+def make_pair(mix_dir, inst_dir):
+    input_exts = ['.wav', '.m4a', '.mp3', '.mp4', '.flac']
+
+    X_list = sorted([
+        os.path.join(mix_dir, fname)
+        for fname in os.listdir(mix_dir)
+        if os.path.splitext(fname)[1] in input_exts
+    ])
+    y_list = sorted([
+        os.path.join(inst_dir, fname)
+        for fname in os.listdir(inst_dir)
+        if os.path.splitext(fname)[1] in input_exts
+    ])
+
+    filelist = list(zip(X_list, y_list))
+
+    return filelist
+
+
+def train_val_split(dataset_dir, split_mode, val_rate, val_filelist):
+    if split_mode == 'random':
+        filelist = make_pair(
+            os.path.join(dataset_dir, 'mixtures'),
+            os.path.join(dataset_dir, 'instruments')
+        )
+
+        random.shuffle(filelist)
+
+        if len(val_filelist) == 0:
+            val_size = int(len(filelist) * val_rate)
+            train_filelist = filelist[:-val_size]
+            val_filelist = filelist[-val_size:]
+        else:
+            train_filelist = [
+                pair for pair in filelist
+                if list(pair) not in val_filelist
+            ]
+    elif split_mode == 'subdirs':
+        if len(val_filelist) != 0:
+            raise ValueError('`val_filelist` option is not available with `subdirs` mode')
+
+        train_filelist = make_pair(
+            os.path.join(dataset_dir, 'training/mixtures'),
+            os.path.join(dataset_dir, 'training/instruments')
+        )
+
+        val_filelist = make_pair(
+            os.path.join(dataset_dir, 'validation/mixtures'),
+            os.path.join(dataset_dir, 'validation/instruments')
+        )
+
+    return train_filelist, val_filelist
+
+
+def make_padding(width, cropsize, offset):
+    left = offset
+    roi_size = cropsize - offset * 2
+    if roi_size == 0:
+        roi_size = cropsize
+    right = roi_size - (width % roi_size) + left
+
+    return left, right, roi_size
+
+
+def make_training_set(filelist, sr, hop_length, n_fft):
+    ret = []
+    for X_path, y_path in tqdm(filelist):
+        X, y, X_cache_path, y_cache_path = spec_utils.cache_or_load(
+            X_path, y_path, sr, hop_length, n_fft
+        )
+        coef = np.max([np.abs(X).max(), np.abs(y).max()])
+        ret.append([X_cache_path, y_cache_path, coef])
+
+    return ret
+
+
+def make_validation_set(filelist, cropsize, sr, hop_length, n_fft, offset):
+    patch_list = []
+    patch_dir = 'cs{}_sr{}_hl{}_nf{}_of{}'.format(cropsize, sr, hop_length, n_fft, offset)
+    os.makedirs(patch_dir, exist_ok=True)
+
+    for X_path, y_path in tqdm(filelist):
+        basename = os.path.splitext(os.path.basename(X_path))[0]
+
+        X, y, _, _ = spec_utils.cache_or_load(X_path, y_path, sr, hop_length, n_fft)
+        coef = np.max([np.abs(X).max(), np.abs(y).max()])
+        X, y = X / coef, y / coef
+
+        l, r, roi_size = make_padding(X.shape[2], cropsize, offset)
+        X_pad = np.pad(X, ((0, 0), (0, 0), (l, r)), mode='constant')
+        y_pad = np.pad(y, ((0, 0), (0, 0), (l, r)), mode='constant')
+
+        len_dataset = int(np.ceil(X.shape[2] / roi_size))
+        for j in range(len_dataset):
+            outpath = os.path.join(patch_dir, '{}_p{}.npz'.format(basename, j))
+            start = j * roi_size
+            if not os.path.exists(outpath):
+                np.savez(
+                    outpath,
+                    X=X_pad[:, :, start:start + cropsize],
+                    y=y_pad[:, :, start:start + cropsize]
+                )
+            patch_list.append(outpath)
+
+    return patch_list
+
+
+def get_oracle_data(X, y, oracle_loss, oracle_rate, oracle_drop_rate):
+    k = int(len(X) * oracle_rate * (1 / (1 - oracle_drop_rate)))
+    n = int(len(X) * oracle_rate)
+    indices = np.argsort(oracle_loss)[::-1][:k]
+    indices = np.random.choice(indices, n, replace=False)
+    oracle_X = X[indices].copy()
+    oracle_y = y[indices].copy()
+
+    return oracle_X, oracle_y, indices
+
+
+if __name__ == "__main__":
+    import sys
+    import utils
+
+    mix_dir = sys.argv[1]
+    inst_dir = sys.argv[2]
+    outdir = sys.argv[3]
+
+    os.makedirs(outdir, exist_ok=True)
+
+    filelist = make_pair(mix_dir, inst_dir)
+    for mix_path, inst_path in tqdm(filelist):
+        mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
+
+        X_spec, y_spec, _, _ = spec_utils.cache_or_load(
+            mix_path, inst_path, 44100, 1024, 2048
+        )
+
+        X_mag = np.abs(X_spec)
+        y_mag = np.abs(y_spec)
+        v_mag = X_mag - y_mag
+        v_mag *= v_mag > y_mag
+
+        outpath = '{}/{}_Vocal.jpg'.format(outdir, mix_basename)
+        v_image = spec_utils.spectrogram_to_image(v_mag)
+        utils.imwrite(outpath, v_image)

lib/layers.py

@@ -0,0 +1,160 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from lib import spec_utils
+
+
+class Conv2DBNActiv(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+        super(Conv2DBNActiv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(
+                nin, nout,
+                kernel_size=ksize,
+                stride=stride,
+                padding=pad,
+                dilation=dilation,
+                bias=False
+            ),
+            nn.BatchNorm2d(nout),
+            activ()
+        )
+
+    def __call__(self, x):
+        return self.conv(x)
+
+
+# class SeperableConv2DBNActiv(nn.Module):
+
+#     def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
+#         super(SeperableConv2DBNActiv, self).__init__()
+#         self.conv = nn.Sequential(
+#             nn.Conv2d(
+#                 nin, nin,
+#                 kernel_size=ksize,
+#                 stride=stride,
+#                 padding=pad,
+#                 dilation=dilation,
+#                 groups=nin,
+#                 bias=False
+#             ),
+#             nn.Conv2d(
+#                 nin, nout,
+#                 kernel_size=1,
+#                 bias=False
+#             ),
+#             nn.BatchNorm2d(nout),
+#             activ()
+#         )
+
+#     def __call__(self, x):
+#         return self.conv(x)
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
+        super(Encoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, stride, pad, activ=activ)
+        self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+
+    def __call__(self, x):
+        h = self.conv1(x)
+        h = self.conv2(h)
+
+        return h
+
+
+class Decoder(nn.Module):
+
+    def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
+        super(Decoder, self).__init__()
+        self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
+        # self.conv2 = Conv2DBNActiv(nout, nout, ksize, 1, pad, activ=activ)
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def __call__(self, x, skip=None):
+        x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)
+
+        if skip is not None:
+            skip = spec_utils.crop_center(skip, x)
+            x = torch.cat([x, skip], dim=1)
+
+        h = self.conv1(x)
+        # h = self.conv2(h)
+
+        if self.dropout is not None:
+            h = self.dropout(h)
+
+        return h
+
+
+class ASPPModule(nn.Module):
+
+    def __init__(self, nin, nout, dilations=(4, 8, 12), activ=nn.ReLU, dropout=False):
+        super(ASPPModule, self).__init__()
+        self.conv1 = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, None)),
+            Conv2DBNActiv(nin, nout, 1, 1, 0, activ=activ)
+        )
+        self.conv2 = Conv2DBNActiv(
+            nin, nout, 1, 1, 0, activ=activ
+        )
+        self.conv3 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[0], dilations[0], activ=activ
+        )
+        self.conv4 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[1], dilations[1], activ=activ
+        )
+        self.conv5 = Conv2DBNActiv(
+            nin, nout, 3, 1, dilations[2], dilations[2], activ=activ
+        )
+        self.bottleneck = Conv2DBNActiv(
+            nout * 5, nout, 1, 1, 0, activ=activ
+        )
+        self.dropout = nn.Dropout2d(0.1) if dropout else None
+
+    def forward(self, x):
+        _, _, h, w = x.size()
+        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
+        feat2 = self.conv2(x)
+        feat3 = self.conv3(x)
+        feat4 = self.conv4(x)
+        feat5 = self.conv5(x)
+        out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
+        out = self.bottleneck(out)
+
+        if self.dropout is not None:
+            out = self.dropout(out)
+
+        return out
+
+
+class LSTMModule(nn.Module):
+
+    def __init__(self, nin_conv, nin_lstm, nout_lstm):
+        super(LSTMModule, self).__init__()
+        self.conv = Conv2DBNActiv(nin_conv, 1, 1, 1, 0)
+        self.lstm = nn.LSTM(
+            input_size=nin_lstm,
+            hidden_size=nout_lstm // 2,
+            bidirectional=True
+        )
+        self.dense = nn.Sequential(
+            nn.Linear(nout_lstm, nin_lstm),
+            nn.BatchNorm1d(nin_lstm),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        N, _, nbins, nframes = x.size()
+        h = self.conv(x)[:, 0]  # N, nbins, nframes
+        h = h.permute(2, 0, 1)  # nframes, N, nbins
+        h, _ = self.lstm(h)
+        h = self.dense(h.reshape(-1, h.size()[-1]))  # nframes * N, nbins
+        h = h.reshape(nframes, N, 1, nbins)
+        h = h.permute(1, 2, 3, 0)
+
+        return h

lib/nets.py

@@ -0,0 +1,131 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from lib import layers
+
+
+class BaseNet(nn.Module):
+
+    def __init__(self, nin, nout, nin_lstm, nout_lstm, dilations=((4, 2), (8, 4), (12, 6))):
+        super(BaseNet, self).__init__()
+        self.enc1 = layers.Conv2DBNActiv(nin, nout, 3, 1, 1)
+        self.enc2 = layers.Encoder(nout, nout * 2, 3, 2, 1)
+        self.enc3 = layers.Encoder(nout * 2, nout * 4, 3, 2, 1)
+        self.enc4 = layers.Encoder(nout * 4, nout * 6, 3, 2, 1)
+        self.enc5 = layers.Encoder(nout * 6, nout * 8, 3, 2, 1)
+
+        self.aspp = layers.ASPPModule(nout * 8, nout * 8, dilations, dropout=True)
+
+        self.dec4 = layers.Decoder(nout * (6 + 8), nout * 6, 3, 1, 1)
+        self.dec3 = layers.Decoder(nout * (4 + 6), nout * 4, 3, 1, 1)
+        self.dec2 = layers.Decoder(nout * (2 + 4), nout * 2, 3, 1, 1)
+        self.lstm_dec2 = layers.LSTMModule(nout * 2, nin_lstm, nout_lstm)
+        self.dec1 = layers.Decoder(nout * (1 + 2) + 1, nout * 1, 3, 1, 1)
+
+    def __call__(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(e1)
+        e3 = self.enc3(e2)
+        e4 = self.enc4(e3)
+        e5 = self.enc5(e4)
+
+        h = self.aspp(e5)
+
+        h = self.dec4(h, e4)
+        h = self.dec3(h, e3)
+        h = self.dec2(h, e2)
+        h = torch.cat([h, self.lstm_dec2(h)], dim=1)
+        h = self.dec1(h, e1)
+
+        return h
+
+
+class CascadedNet(nn.Module):
+
+    def __init__(self, n_fft, nout=32, nout_lstm=128):
+        super(CascadedNet, self).__init__()
+        self.max_bin = n_fft // 2
+        self.output_bin = n_fft // 2 + 1
+        self.nin_lstm = self.max_bin // 2
+        self.offset = 64
+
+        self.stg1_low_band_net = nn.Sequential(
+            BaseNet(2, nout // 2, self.nin_lstm // 2, nout_lstm),
+            layers.Conv2DBNActiv(nout // 2, nout // 4, 1, 1, 0)
+        )
+        self.stg1_high_band_net = BaseNet(
+            2, nout // 4, self.nin_lstm // 2, nout_lstm // 2
+        )
+
+        self.stg2_low_band_net = nn.Sequential(
+            BaseNet(nout // 4 + 2, nout, self.nin_lstm // 2, nout_lstm),
+            layers.Conv2DBNActiv(nout, nout // 2, 1, 1, 0)
+        )
+        self.stg2_high_band_net = BaseNet(
+            nout // 4 + 2, nout // 2, self.nin_lstm // 2, nout_lstm // 2
+        )
+
+        self.stg3_full_band_net = BaseNet(
+            3 * nout // 4 + 2, nout, self.nin_lstm, nout_lstm
+        )
+
+        self.out = nn.Conv2d(nout, 2, 1, bias=False)
+        self.aux_out = nn.Conv2d(3 * nout // 4, 2, 1, bias=False)
+
+    def forward(self, x):
+        x = x[:, :, :self.max_bin]
+
+        bandw = x.size()[2] // 2
+        l1_in = x[:, :, :bandw]
+        h1_in = x[:, :, bandw:]
+        l1 = self.stg1_low_band_net(l1_in)
+        h1 = self.stg1_high_band_net(h1_in)
+        aux1 = torch.cat([l1, h1], dim=2)
+
+        l2_in = torch.cat([l1_in, l1], dim=1)
+        h2_in = torch.cat([h1_in, h1], dim=1)
+        l2 = self.stg2_low_band_net(l2_in)
+        h2 = self.stg2_high_band_net(h2_in)
+        aux2 = torch.cat([l2, h2], dim=2)
+
+        f3_in = torch.cat([x, aux1, aux2], dim=1)
+        f3 = self.stg3_full_band_net(f3_in)
+
+        mask = torch.sigmoid(self.out(f3))
+        mask = F.pad(
+            input=mask,
+            pad=(0, 0, 0, self.output_bin - mask.size()[2]),
+            mode='replicate'
+        )
+
+        if self.training:
+            aux = torch.cat([aux1, aux2], dim=1)
+            aux = torch.sigmoid(self.aux_out(aux))
+            aux = F.pad(
+                input=aux,
+                pad=(0, 0, 0, self.output_bin - aux.size()[2]),
+                mode='replicate'
+            )
+            return mask, aux
+        else:
+            return mask
+
+    def predict_mask(self, x):
+        mask = self.forward(x)
+
+        if self.offset > 0:
+            mask = mask[:, :, :, self.offset:-self.offset]
+            assert mask.size()[3] > 0
+
+        return mask
+
+    def predict(self, x):
+        mask = self.forward(x)
+        pred_mag = x * mask
+
+        if self.offset > 0:
+            pred_mag = pred_mag[:, :, :, self.offset:-self.offset]
+            assert pred_mag.size()[3] > 0
+
+        return pred_mag

lib/spec_utils.py

@@ -0,0 +1,227 @@
+import os
+
+import librosa
+import numpy as np
+import soundfile as sf
+
+
+def crop_center(h1, h2):
+    h1_shape = h1.size()
+    h2_shape = h2.size()
+
+    if h1_shape[3] == h2_shape[3]:
+        return h1
+    elif h1_shape[3] < h2_shape[3]:
+        raise ValueError('h1_shape[3] must be greater than h2_shape[3]')
+
+    # s_freq = (h2_shape[2] - h1_shape[2]) // 2
+    # e_freq = s_freq + h1_shape[2]
+    s_time = (h1_shape[3] - h2_shape[3]) // 2
+    e_time = s_time + h2_shape[3]
+    h1 = h1[:, :, :, s_time:e_time]
+
+    return h1
+
+
+def wave_to_spectrogram(wave, hop_length, n_fft):
+    wave_left = np.asfortranarray(wave[0])
+    wave_right = np.asfortranarray(wave[1])
+
+    spec_left = librosa.stft(wave_left, n_fft, hop_length=hop_length)
+    spec_right = librosa.stft(wave_right, n_fft, hop_length=hop_length)
+    spec = np.asfortranarray([spec_left, spec_right])
+
+    return spec
+
+
+def spectrogram_to_image(spec, mode='magnitude'):
+    if mode == 'magnitude':
+        if np.iscomplexobj(spec):
+            y = np.abs(spec)
+        else:
+            y = spec
+        y = np.log10(y ** 2 + 1e-8)
+    elif mode == 'phase':
+        if np.iscomplexobj(spec):
+            y = np.angle(spec)
+        else:
+            y = spec
+
+    y -= y.min()
+    y *= 255 / y.max()
+    img = np.uint8(y)
+
+    if y.ndim == 3:
+        img = img.transpose(1, 2, 0)
+        img = np.concatenate([
+            np.max(img, axis=2, keepdims=True), img
+        ], axis=2)
+
+    return img
+
+
+def aggressively_remove_vocal(X, y, weight):
+    X_mag = np.abs(X)
+    y_mag = np.abs(y)
+    # v_mag = np.abs(X_mag - y_mag)
+    v_mag = X_mag - y_mag
+    v_mag *= v_mag > y_mag
+
+    y_mag = np.clip(y_mag - v_mag * weight, 0, np.inf)
+
+    return y_mag * np.exp(1.j * np.angle(y))
+
+
+def merge_artifacts(y_mask, thres=0.05, min_range=64, fade_size=32):
+    if min_range < fade_size * 2:
+        raise ValueError('min_range must be >= fade_size * 2')
+
+    idx = np.where(y_mask.min(axis=(0, 1)) > thres)[0]
+    start_idx = np.insert(idx[np.where(np.diff(idx) != 1)[0] + 1], 0, idx[0])
+    end_idx = np.append(idx[np.where(np.diff(idx) != 1)[0]], idx[-1])
+    artifact_idx = np.where(end_idx - start_idx > min_range)[0]
+    weight = np.zeros_like(y_mask)
+    if len(artifact_idx) > 0:
+        start_idx = start_idx[artifact_idx]
+        end_idx = end_idx[artifact_idx]
+        old_e = None
+        for s, e in zip(start_idx, end_idx):
+            if old_e is not None and s - old_e < fade_size:
+                s = old_e - fade_size * 2
+
+            if s != 0:
+                weight[:, :, s:s + fade_size] = np.linspace(0, 1, fade_size)
+            else:
+                s -= fade_size
+
+            if e != y_mask.shape[2]:
+                weight[:, :, e - fade_size:e] = np.linspace(1, 0, fade_size)
+            else:
+                e += fade_size
+
+            weight[:, :, s + fade_size:e - fade_size] = 1
+            old_e = e
+
+    v_mask = 1 - y_mask
+    y_mask += weight * v_mask
+
+    return y_mask
+
+
+def align_wave_head_and_tail(a, b, sr):
+    a, _ = librosa.effects.trim(a)
+    b, _ = librosa.effects.trim(b)
+
+    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, n_fft):
+    mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
+    inst_basename = os.path.splitext(os.path.basename(inst_path))[0]
+
+    cache_dir = 'sr{}_hl{}_nf{}'.format(sr, hop_length, n_fft)
+    mix_cache_dir = os.path.join(os.path.dirname(mix_path), cache_dir)
+    inst_cache_dir = os.path.join(os.path.dirname(inst_path), cache_dir)
+    os.makedirs(mix_cache_dir, exist_ok=True)
+    os.makedirs(inst_cache_dir, exist_ok=True)
+
+    mix_cache_path = os.path.join(mix_cache_dir, mix_basename + '.npy')
+    inst_cache_path = os.path.join(inst_cache_dir, inst_basename + '.npy')
+
+    if os.path.exists(mix_cache_path) and os.path.exists(inst_cache_path):
+        X = np.load(mix_cache_path)
+        y = np.load(inst_cache_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, y = align_wave_head_and_tail(X, y, sr)
+
+        X = wave_to_spectrogram(X, hop_length, n_fft)
+        y = wave_to_spectrogram(y, hop_length, n_fft)
+
+        np.save(mix_cache_path, X)
+        np.save(inst_cache_path, y)
+
+    return X, y, mix_cache_path, inst_cache_path
+
+
+def spectrogram_to_wave(spec, hop_length=1024):
+    if spec.ndim == 2:
+        wave = librosa.istft(spec, hop_length=hop_length)
+    elif spec.ndim == 3:
+        spec_left = np.asfortranarray(spec[0])
+        spec_right = np.asfortranarray(spec[1])
+
+        wave_left = librosa.istft(spec_left, hop_length=hop_length)
+        wave_right = librosa.istft(spec_right, hop_length=hop_length)
+        wave = np.asfortranarray([wave_left, wave_right])
+
+    return wave
+
+
+if __name__ == "__main__":
+    import cv2
+    import sys
+
+    bins = 2048 // 2 + 1
+    freq_to_bin = 2 * bins / 44100
+    unstable_bins = int(200 * freq_to_bin)
+    stable_bins = int(22050 * freq_to_bin)
+    reduction_weight = np.concatenate([
+        np.linspace(0, 1, unstable_bins, dtype=np.float32)[:, None],
+        np.linspace(1, 0, stable_bins - unstable_bins, dtype=np.float32)[:, None],
+        np.zeros((bins - stable_bins, 1))
+    ], axis=0) * 0.2
+
+    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, y = align_wave_head_and_tail(X, y, 44100)
+    X_spec = wave_to_spectrogram(X, 1024, 2048)
+    y_spec = wave_to_spectrogram(y, 1024, 2048)
+
+    X_mag = np.abs(X_spec)
+    y_mag = np.abs(y_spec)
+    # v_mag = np.abs(X_mag - y_mag)
+    v_mag = X_mag - y_mag
+    v_mag *= v_mag > y_mag
+
+    # y_mag = np.clip(y_mag - v_mag * reduction_weight, 0, np.inf)
+    y_spec = y_mag * np.exp(1j * np.angle(y_spec))
+    v_spec = v_mag * np.exp(1j * np.angle(X_spec))
+
+    X_image = spectrogram_to_image(X_mag)
+    y_image = spectrogram_to_image(y_mag)
+    v_image = spectrogram_to_image(v_mag)
+
+    cv2.imwrite('test_X.jpg', X_image)
+    cv2.imwrite('test_y.jpg', y_image)
+    cv2.imwrite('test_v.jpg', v_image)
+
+    sf.write('test_X.wav', spectrogram_to_wave(X_spec).T, 44100)
+    sf.write('test_y.wav', spectrogram_to_wave(y_spec).T, 44100)
+    sf.write('test_v.wav', spectrogram_to_wave(v_spec).T, 44100)

lib/utils.py

@@ -0,0 +1,30 @@
+import os
+
+import cv2
+import numpy as np
+
+
+def imread(filename, flags=cv2.IMREAD_COLOR, dtype=np.uint8):
+    try:
+        n = np.fromfile(filename, dtype)
+        img = cv2.imdecode(n, flags)
+        return img
+    except Exception as e:
+        print(e)
+        return None
+
+
+def imwrite(filename, img, params=None):
+    try:
+        ext = os.path.splitext(filename)[1]
+        result, n = cv2.imencode(ext, img, params)
+
+        if result:
+            with open(filename, mode='w+b') as f:
+                n.tofile(f)
+            return True
+        else:
+            return False
+    except Exception as e:
+        print(e)
+        return False

requirements.txt

@@ -0,0 +1,5 @@
+torch
+tqdm>=4.30
+librosa>=0.6.3,<0.9
+opencv_python>=4.2.0
+hopsworks