97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
import pytest
|
|
from torchvision.transforms import ToPILImage
|
|
from .conftest import TestData
|
|
from ..core.plotting import get_spectogram, get_wave, save_image
|
|
from ..visuals.SignalProcessingSpectrogram import SignalProcessingSpectrogram
|
|
|
|
TEST_NAME = "plotting"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"n_fft, n_mels, xlim",
|
|
[
|
|
(4096, 128 * 1, 8192),
|
|
(4096, 128 * 2, 8192),
|
|
(4096, 128 * 4, 8192),
|
|
(4096, 128 * 6, 8192),
|
|
],
|
|
)
|
|
def test_plotting_spectogram_general(
|
|
test_data: TestData, n_fft: int, n_mels: int, xlim: int
|
|
) -> None:
|
|
|
|
waveform = test_data["audio"]["waveform"].squeeze(0)
|
|
sample_rate = test_data["audio"]["sample_rate"]
|
|
|
|
print("waveform", waveform.shape)
|
|
|
|
spectogram = get_spectogram(
|
|
waveform, sample_rate=sample_rate, n_fft=n_fft, n_mels=n_mels, xlim=xlim
|
|
)
|
|
|
|
pest_param_str = test_data["pest_param_str"]
|
|
output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
|
|
|
|
save_image(output_filepath, spectogram)
|
|
|
|
assert output_filepath.exists(), f"Output file {output_filepath} was not created."
|
|
assert output_filepath is not None, "Processed audio output is None."
|
|
|
|
print(f"test_plotting_spectogram_general {test_data['output_root']}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"stretch_factor, window_size_seconds",
|
|
[(8.0, 0.25)],
|
|
)
|
|
def test_plotting_waveform_general(
|
|
test_data: TestData, stretch_factor: float, window_size_seconds: float
|
|
) -> None:
|
|
|
|
waveform = test_data["audio"]["waveform"].squeeze(0)
|
|
sample_rate = test_data["audio"]["sample_rate"]
|
|
|
|
print("waveform", waveform.shape)
|
|
|
|
spectogram = get_wave(waveform, sample_rate=sample_rate, xlim=4096)
|
|
|
|
pest_param_str = test_data["pest_param_str"]
|
|
output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
|
|
|
|
save_image(output_filepath, spectogram)
|
|
|
|
assert output_filepath.exists(), f"Output file {output_filepath} was not created."
|
|
assert output_filepath is not None, "Processed audio output is None."
|
|
|
|
print(f"test_plotting_spectogram_general {test_data['output_root']}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"stretch_factor, window_size_seconds",
|
|
[(8.0, 0.25)],
|
|
)
|
|
def test_plotting_waveform_node_general(
|
|
test_data: TestData, stretch_factor: float, window_size_seconds: float
|
|
) -> None:
|
|
|
|
node = SignalProcessingSpectrogram()
|
|
|
|
output = node.process(audio_input=test_data["audio"])[0]
|
|
|
|
pest_param_str = test_data["pest_param_str"]
|
|
output_filepath = test_data["output_root"] / f"{pest_param_str}.png"
|
|
|
|
to_pil = ToPILImage()
|
|
print("rgb_image[0] -----------------------------", output[0].shape)
|
|
rgb_image = output[0][..., :3]
|
|
rgb_image = rgb_image.permute(2, 0, 1)
|
|
print("rgb_image -----------------------------", rgb_image.shape)
|
|
spectogram = to_pil(rgb_image)
|
|
print("spectogram -----------------------------", spectogram)
|
|
save_image(output_filepath, spectogram)
|
|
|
|
assert output_filepath.exists(), f"Output file {output_filepath} was not created."
|
|
assert output is not None, "Processed audio output is None."
|
|
|
|
print(f"test_pitch_shift_general {test_data['output_root']}")
|