feat: deterministic shifts in input samples

requirements.txt

@@ -8,3 +8,4 @@ musdb
 SoundFile
 scipy
 norbert
+tqdm

test_demucs.py

@@ -3,22 +3,28 @@
 import numpy as np
 import soundfile as sf
 import torch
-
-
+from tqdm import tqdm
 from evaluator.music_demixing import MusicDemixingPredictor
 
 
 def separate(
     audio,
+    rate,
     model,
+    shifts=10,
     device="cpu"
 ):
 
     # convert numpy audio to torch
     audio_torch = torch.tensor(audio.T).float().to(device)
+    max_shift = int(rate * 0.5)
 
     with torch.no_grad():
-        y = model(audio_torch.unsqueeze(0)).squeeze()
+        y = torch.zeros_like(audio_torch).repeat(4, 1, 1)
+        for shift in tqdm(torch.linspace(0, max_shift, steps=shifts, dtype=torch.long).tolist()):
+            shifted_y = model(audio_torch.unsqueeze(0)[..., shift:]).squeeze()
+            y[..., shift:] += shifted_y
+        y /= shifts
 
     return y.numpy()
 
@@ -45,6 +51,7 @@ class DemucsPredictor(MusicDemixingPredictor):
         # Step 3: Perform separation
         estimates = separate(
             x,
+            rate,
             self.separator,
         )