mirror of
https://github.com/YuzuZensai/spleeter.git
synced 2026-01-06 04:32:43 +00:00
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf8
|
|
|
|
""" Unit testing for Separator class. """
|
|
|
|
__email__ = 'research@deezer.com'
|
|
__author__ = 'Deezer Research'
|
|
__license__ = 'MIT License'
|
|
|
|
from os.path import exists, join
|
|
from tempfile import TemporaryDirectory
|
|
|
|
import pytest
|
|
|
|
from spleeter.audio.adapter import get_default_audio_adapter
|
|
from spleeter.separator import Separator
|
|
|
|
TEST_AUDIO_DESCRIPTOR = 'audio_example.mp3'
|
|
TEST_CONFIGURATIONS = [
|
|
('spleeter:2stems', ('vocals', 'accompaniment')),
|
|
('spleeter:4stems', ('vocals', 'drums', 'bass', 'other')),
|
|
('spleeter:5stems', ('vocals', 'drums', 'bass', 'piano', 'other'))
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize('configuration, instruments', TEST_CONFIGURATIONS)
|
|
def test_separate(configuration, instruments):
|
|
""" Test separation from raw data. """
|
|
adapter = get_default_audio_adapter()
|
|
waveform, _ = adapter.load(TEST_AUDIO_DESCRIPTOR)
|
|
separator = Separator(configuration)
|
|
prediction = separator.separate(waveform)
|
|
assert len(prediction) == 2
|
|
for instrument in instruments:
|
|
assert instrument in prediction
|
|
|
|
|
|
@pytest.mark.parametrize('configuration, instruments', TEST_CONFIGURATIONS)
|
|
def test_separate_to_file(configuration, instruments):
|
|
""" Test file based separation. """
|
|
separator = Separator(configuration)
|
|
with TemporaryDirectory() as directory:
|
|
separator.separate_to_file(
|
|
TEST_AUDIO_DESCRIPTOR,
|
|
directory)
|
|
for instrument in instruments:
|
|
assert exists(join(directory, '{}.wav'.format(instrument)))
|
|
# TODO: Consider testing generated file as well.
|