Adding support for mono files with librosa backend

spleeter/separator.py

@@ -123,12 +123,16 @@ def stft(self, data, inverse=False, length=None):
         win = hann(N, sym=False)
         fstft = istft if inverse else stft
         win_len_arg = {"win_length": None, "length": length} if inverse else {"n_fft": N}
-        dl, dr = (data[:, :, 0].T, data[:, :, 1].T) if inverse else (data[:, 0], data[:, 1])
-        s1 = fstft(dl, hop_length=H, window=win, center=False, **win_len_arg)
-        s2 = fstft(dr, hop_length=H, window=win, center=False, **win_len_arg)
-        s1 = np.expand_dims(s1.T, 2-inverse)
-        s2 = np.expand_dims(s2.T, 2-inverse)
-        return np.concatenate([s1, s2], axis=2-inverse)
+        n_channels = data.shape[-1]
+        out = []
+        for c in range(n_channels):
+            d = data[:, :, c].T if inverse else data[:, c]
+            s = fstft(dl, hop_length=H, window=win, center=False, **win_len_arg)
+            s = np.expand_dims(s.T, 2-inverse)
+            out.append(s)
+        if len(out) == 1:
+            return out[0]
+        return np.concatenate(out, axis=2-inverse)
 
     def separate_librosa(self, waveform, audio_id):
         out = {}

tests/test_separator.py

@@ -8,7 +8,7 @@
 __license__ = 'MIT License'
 
 import filecmp
-
+import itertools
 from os.path import splitext, basename, exists, join
 from tempfile import TemporaryDirectory
 
@@ -19,25 +19,26 @@
 from spleeter.audio.adapter import get_default_audio_adapter
 from spleeter.separator import Separator
 
-TEST_AUDIO_DESCRIPTOR = 'audio_example.mp3'
-TEST_AUDIO_BASENAME = splitext(basename(TEST_AUDIO_DESCRIPTOR))[0]
-TEST_CONFIGURATIONS = [
-    ('spleeter:2stems', ('vocals', 'accompaniment'), 'tensorflow'),
-    ('spleeter:4stems', ('vocals', 'drums', 'bass', 'other'), 'tensorflow'),
-    ('spleeter:5stems', ('vocals', 'drums', 'bass', 'piano', 'other'), 'tensorflow'),
-    ('spleeter:2stems', ('vocals', 'accompaniment'), 'librosa'),
-    ('spleeter:4stems', ('vocals', 'drums', 'bass', 'other'), 'librosa'),
-    ('spleeter:5stems', ('vocals', 'drums', 'bass', 'piano', 'other'), 'librosa')
-]
+TEST_AUDIO_DESCRIPTORS = ['audio_example.mp3', 'audio_example_mono.mp3']
+BACKENDS = ["tensorflow", "librosa"]
+MODELS = ['spleeter:2stems', 'spleeter:4stems', 'spleeter:5stems']
+MODEL_TO_INST = {
+    'spleeter:2stems': ('vocals', 'accompaniment'),
+    'spleeter:4stems': ('vocals', 'drums', 'bass', 'other'),
+    'spleeter:5stems': ('vocals', 'drums', 'bass', 'piano', 'other'),
+}
+
+
+TEST_CONFIGURATIONS = list(itertools.product(TEST_AUDIO_DESCRIPTORS, MODELS, BACKENDS))
 
 
 @pytest.mark.parametrize('configuration, instruments, backend', TEST_CONFIGURATIONS)
-def test_separate(configuration, instruments, backend):
+def test_separate(test_file, configuration, instruments, backend):
     """ Test separation from raw data. """
     adapter = get_default_audio_adapter()
-    waveform, _ = adapter.load(TEST_AUDIO_DESCRIPTOR)
+    waveform, _ = adapter.load(test_file)
     separator = Separator(configuration, stft_backend=backend)
-    prediction = separator.separate(waveform, TEST_AUDIO_DESCRIPTOR)
+    prediction = separator.separate(waveform, test_file)
     assert len(prediction) == len(instruments)
     for instrument in instruments:
         assert instrument in prediction
@@ -51,40 +52,42 @@ def test_separate(configuration, instruments, backend):
 
 
 @pytest.mark.parametrize('configuration, instruments, backend', TEST_CONFIGURATIONS)
-def test_separate_to_file(configuration, instruments, backend):
+def test_separate_to_file(test_file, configuration, instruments, backend):
     """ Test file based separation. """
     separator = Separator(configuration, stft_backend=backend)
+    basename = splitext(basename(test_file))
     with TemporaryDirectory() as directory:
         separator.separate_to_file(
-            TEST_AUDIO_DESCRIPTOR,
+            test_file,
             directory)
         for instrument in instruments:
             assert exists(join(
                 directory,
-                '{}/{}.wav'.format(TEST_AUDIO_BASENAME, instrument)))
+                '{}/{}.wav'.format(basename, instrument)))
 
 
 @pytest.mark.parametrize('configuration, instruments, backend', TEST_CONFIGURATIONS)
-def test_filename_format(configuration, instruments, backend):
+def test_filename_format(test_file, configuration, instruments, backend):
     """ Test custom filename format. """
     separator = Separator(configuration, stft_backend=backend)
+    basename = splitext(basename(test_file))
     with TemporaryDirectory() as directory:
         separator.separate_to_file(
-            TEST_AUDIO_DESCRIPTOR,
+            test_file,
             directory,
             filename_format='export/{filename}/{instrument}.{codec}')
         for instrument in instruments:
             assert exists(join(
                 directory,
-                'export/{}/{}.wav'.format(TEST_AUDIO_BASENAME, instrument)))
+                'export/{}/{}.wav'.format(basename, instrument)))
 
 
-def test_filename_conflict():
+def test_filename_conflict(test_file):
     """ Test error handling with static pattern. """
     separator = Separator(TEST_CONFIGURATIONS[0][0])
     with TemporaryDirectory() as directory:
         with pytest.raises(SpleeterError):
             separator.separate_to_file(
-                TEST_AUDIO_DESCRIPTOR,
+                test_file,
                 directory,
                 filename_format='I wanna be your lover')