Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,997 Bytes
01f8b5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
import os
import pytest
import tempfile
import shutil
from pathlib import Path
from audio_separator.separator import Separator
@pytest.fixture(name="input_file")
def fixture_input_file():
"""Fixture providing the test input audio file path."""
return "tests/inputs/mardy20s.flac"
@pytest.fixture(name="temp_output_dir")
def fixture_temp_output_dir():
"""Fixture providing a temporary directory for output files."""
temp_dir = tempfile.mkdtemp()
yield temp_dir
# Clean up after test
shutil.rmtree(temp_dir)
def test_separator_output_dir_and_custom_names(input_file, temp_output_dir):
"""Test that Separator respects output_dir and custom_output_names parameters."""
print("\n>>> TEST: Checking output_dir with custom output names")
# Define custom output filenames
vocal_output_filename = "custom_vocals_output"
instrumental_output_filename = "custom_instrumental_output"
# Create output name mapping
custom_output_names = {"Vocals": vocal_output_filename, "Instrumental": instrumental_output_filename}
# Initialize separator with specified output directory
print(f"Creating Separator with output_dir: {temp_output_dir}")
separator = Separator(output_dir=temp_output_dir, log_level=20) # INFO level
# Load model
print("Loading model: MGM_MAIN_v4.pth")
separator.load_model(model_filename="MGM_MAIN_v4.pth")
# Run separation with custom output names
print(f"Running separation with custom_output_names: {custom_output_names}")
output_files = separator.separate(input_file, custom_output_names=custom_output_names)
print(f"Separator.separate() returned: {output_files}")
# The separator adds .wav extension since the default output format is WAV
expected_vocal_filename = vocal_output_filename + ".wav"
expected_instrumental_filename = instrumental_output_filename + ".wav"
# Check that the returned filenames match the expected names
output_filenames = [os.path.basename(f) for f in output_files]
print(f"Extracted filenames from output_files: {output_filenames}")
# NOTE: The Separator class returns only the filenames, not the full paths with output_dir
print("EXPECTED BEHAVIOR: The Separator.separate() method returns filenames without the output_dir path")
print(f"Expected filenames (without path): {expected_vocal_filename} and {expected_instrumental_filename}")
assert expected_vocal_filename in output_filenames, f"Expected {expected_vocal_filename} in output files"
assert expected_instrumental_filename in output_filenames, f"Expected {expected_instrumental_filename} in output files"
# Check that files physically exist in the specified output directory
expected_vocal_path = os.path.join(temp_output_dir, expected_vocal_filename)
expected_instrumental_path = os.path.join(temp_output_dir, expected_instrumental_filename)
print(f"Checking that files exist in output_dir: {temp_output_dir}")
print(f"Full expected vocal path: {expected_vocal_path}")
print(f"Full expected instrumental path: {expected_instrumental_path}")
assert os.path.exists(expected_vocal_path), f"Vocals output file doesn't exist: {expected_vocal_path}"
assert os.path.exists(expected_instrumental_path), f"Instrumental output file doesn't exist: {expected_instrumental_path}"
assert os.path.getsize(expected_vocal_path) > 0, f"Vocals output file is empty: {expected_vocal_path}"
assert os.path.getsize(expected_instrumental_path) > 0, f"Instrumental output file is empty: {expected_instrumental_path}"
print("✅ Test passed: Separator correctly handles output_dir and custom_output_names")
print(" - Files were saved to the specified output directory")
print(" - Custom filenames were used (with .wav extension added)")
print(" - Returned paths include just the filenames (not the full paths)")
def test_separator_single_stem_output(input_file, temp_output_dir):
"""Test that Separator correctly respects output_single_stem with custom output name."""
print("\n>>> TEST: Checking output_single_stem with custom output name")
# Define custom output filename for single stem
vocal_output_filename = "only_vocals_output"
# Create output name mapping
custom_output_names = {"Vocals": vocal_output_filename}
# Initialize separator with specified output directory and single stem output
print(f"Creating Separator with output_dir: {temp_output_dir} and output_single_stem: Vocals")
separator = Separator(
output_dir=temp_output_dir,
output_single_stem="Vocals", # Only extract vocals
log_level=20 # INFO level
)
# Load model
print("Loading model: MGM_MAIN_v4.pth")
separator.load_model(model_filename="MGM_MAIN_v4.pth")
# Run separation with custom output name
print(f"Running separation with custom_output_names: {custom_output_names}")
output_files = separator.separate(input_file, custom_output_names=custom_output_names)
print(f"Separator.separate() returned: {output_files}")
# The separator adds .wav extension since the default output format is WAV
expected_vocal_filename = vocal_output_filename + ".wav"
# Check that the output files list contains only the expected file
print(f"Checking that only one file was returned and it has the correct name")
assert len(output_files) == 1, f"Expected 1 output file, got {len(output_files)}"
assert os.path.basename(output_files[0]) == expected_vocal_filename, f"Expected {expected_vocal_filename} in output files"
# Check that file physically exists in the specified output directory
expected_vocal_path = os.path.join(temp_output_dir, expected_vocal_filename)
print(f"Checking that file exists in output_dir: {expected_vocal_path}")
assert os.path.exists(expected_vocal_path), f"Vocals output file doesn't exist: {expected_vocal_path}"
assert os.path.getsize(expected_vocal_path) > 0, f"Vocals output file is empty: {expected_vocal_path}"
# Make sure the instrumental file was NOT created
instrumental_files = [f for f in os.listdir(temp_output_dir) if "instrumental" in f.lower()]
print(f"Checking that no instrumental files were created, found: {instrumental_files}")
assert len(instrumental_files) == 0, f"No instrumental file should be created when using output_single_stem, found: {instrumental_files}"
print("✅ Test passed: Separator correctly handles output_single_stem with custom output name")
print(" - Only the specified stem (Vocals) was extracted")
print(" - The custom filename was used (with .wav extension added)")
print(" - No other stems (like Instrumental) were created")
def test_separator_output_without_custom_names(input_file, temp_output_dir):
"""Test that Separator respects output_dir without custom_output_names."""
print("\n>>> TEST: Checking output_dir without custom output names")
# Initialize separator with specified output directory
print(f"Creating Separator with output_dir: {temp_output_dir}")
separator = Separator(output_dir=temp_output_dir, log_level=20) # INFO level
# Load model
print("Loading model: MGM_MAIN_v4.pth")
separator.load_model(model_filename="MGM_MAIN_v4.pth")
# Run separation without custom output names
print("Running separation without custom_output_names")
output_files = separator.separate(input_file)
print(f"Separator.separate() returned: {output_files}")
# Check that output files exist and have content
print(f"Checking that two files were returned (one for each stem)")
assert len(output_files) == 2, f"Expected 2 output files, got {len(output_files)}"
# Check if the files were created in the output directory
# Note: The separator doesn't include the full path in the returned output_files
output_file_basenames = [os.path.basename(f) for f in output_files]
print(f"Extracted filenames from output_files: {output_file_basenames}")
print("EXPECTED BEHAVIOR: Default filenames are being used (format: inputname_(StemName)_modelname.wav)")
# Check all output files exist in the specified output directory
print(f"Checking that files exist in output_dir: {temp_output_dir}")
for basename in output_file_basenames:
full_path = os.path.join(temp_output_dir, basename)
print(f"Checking file: {full_path}")
assert os.path.exists(full_path), f"Output file doesn't exist: {full_path}"
assert os.path.getsize(full_path) > 0, f"Output file is empty: {full_path}"
print("✅ Test passed: Separator correctly handles output_dir without custom output names")
print(" - Files were saved to the specified output directory")
print(" - Default naming scheme was used (input_name_(Stem)_model.wav)")
print(" - Returned paths include just the filenames (not the full paths)")
|