diff --git a/audio_example_mono.mp3 b/audio_example_mono.mp3 new file mode 100644 index 0000000..7b134ac Binary files /dev/null and b/audio_example_mono.mp3 differ diff --git a/spleeter/separator.py b/spleeter/separator.py index 174c73f..9b77111 100644 --- a/spleeter/separator.py +++ b/spleeter/separator.py @@ -123,12 +123,16 @@ class Separator(object): 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 = {} diff --git a/tests/test_separator.py b/tests/test_separator.py index 271fdfb..85643c1 100644 --- a/tests/test_separator.py +++ b/tests/test_separator.py @@ -8,7 +8,7 @@ __author__ = 'Deezer Research' __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 import SpleeterError 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')