251 lines
8.7 KiB
Python
251 lines
8.7 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from src.doa_estimator import (
|
|
DOA_ALGORITHMS,
|
|
DoaConfig,
|
|
DoaEstimator,
|
|
load_pyroomacoustics_normmusic,
|
|
load_pyroomacoustics_srp,
|
|
)
|
|
|
|
|
|
def circular_error_degrees(actual: float, expected: float) -> float:
|
|
return abs((actual - expected + 180.0) % 360.0 - 180.0)
|
|
|
|
|
|
def synthesize_plane_wave(
|
|
estimator: DoaEstimator,
|
|
angle_degrees: float,
|
|
) -> np.ndarray:
|
|
config = estimator.config
|
|
rng = np.random.default_rng(8)
|
|
source = rng.standard_normal(config.window_samples)
|
|
frequencies = np.fft.rfftfreq(config.window_samples, 1.0 / config.sample_rate)
|
|
source_spectrum = np.fft.rfft(source)
|
|
|
|
angle = np.deg2rad(angle_degrees)
|
|
direction = np.array([np.cos(angle), np.sin(angle)])
|
|
relative_times = (
|
|
direction @ estimator.microphone_positions / config.sound_speed_m_s
|
|
)
|
|
channels = [
|
|
np.fft.irfft(
|
|
source_spectrum * np.exp(2j * np.pi * frequencies * relative_time),
|
|
n=config.window_samples,
|
|
)
|
|
for relative_time in relative_times
|
|
]
|
|
return (0.05 * np.stack(channels, axis=1)).astype(np.float32)
|
|
|
|
|
|
def test_loads_srp_from_cloned_source() -> None:
|
|
srp = load_pyroomacoustics_srp()
|
|
assert srp.__name__ == "SRP"
|
|
assert srp.__module__.endswith(".srp")
|
|
|
|
|
|
def test_loads_normmusic_from_cloned_source() -> None:
|
|
normmusic = load_pyroomacoustics_normmusic()
|
|
assert normmusic.__name__ == "NormMUSIC"
|
|
assert normmusic.__module__.endswith(".normmusic")
|
|
|
|
|
|
@pytest.mark.parametrize("algorithm", DOA_ALGORITHMS)
|
|
@pytest.mark.parametrize("expected", [0.0, 47.0, 123.0, 275.0, 359.0])
|
|
def test_plane_wave_direction(algorithm: str, expected: float) -> None:
|
|
estimator = DoaEstimator(algorithm=algorithm, start_worker=False)
|
|
try:
|
|
result = estimator.estimate(synthesize_plane_wave(estimator, expected))
|
|
finally:
|
|
estimator.close()
|
|
|
|
assert result.angle_degrees is not None
|
|
assert result.algorithm == algorithm
|
|
assert circular_error_degrees(result.angle_degrees, expected) <= 1.0
|
|
assert result.peak_sharpness > 0.8
|
|
|
|
|
|
def test_runtime_algorithm_switch_clears_previous_result() -> None:
|
|
estimator = DoaEstimator(start_worker=False)
|
|
try:
|
|
assert estimator.algorithm == "SRP-PHAT"
|
|
assert estimator.toggle_algorithm() == "NormMUSIC"
|
|
assert estimator.algorithm == "NormMUSIC"
|
|
assert estimator._engines["NormMUSIC"].frequency_normalization is True
|
|
assert estimator.set_algorithm("MUSIC") == "NormMUSIC"
|
|
assert estimator.set_algorithm("SRP") == "SRP-PHAT"
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_manual_mvdr_steering_angle_is_normalized_and_optional() -> None:
|
|
estimator = DoaEstimator(start_worker=False)
|
|
try:
|
|
assert estimator.mvdr_steering_angle_override is None
|
|
assert estimator.set_mvdr_steering_angle(370.0) == 10.0
|
|
assert estimator.mvdr_steering_angle_override == 10.0
|
|
assert estimator.set_mvdr_steering_angle(-15.0) == 345.0
|
|
assert estimator.set_mvdr_steering_angle(None) is None
|
|
assert estimator.mvdr_steering_angle_override is None
|
|
with pytest.raises(ValueError):
|
|
estimator.set_mvdr_steering_angle(float("nan"))
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_stft_defaults_preserve_48khz_time_and_frequency_resolution() -> None:
|
|
config_48k = DoaConfig(sample_rate=48_000)
|
|
config_24k = DoaConfig(sample_rate=24_000)
|
|
|
|
assert (config_48k.nfft, config_48k.hop, config_48k.window_samples) == (
|
|
512,
|
|
256,
|
|
8_192,
|
|
)
|
|
assert (config_24k.nfft, config_24k.hop, config_24k.window_samples) == (
|
|
256,
|
|
128,
|
|
4_096,
|
|
)
|
|
assert config_48k.sample_rate / config_48k.nfft == (
|
|
config_24k.sample_rate / config_24k.nfft
|
|
)
|
|
assert config_48k.window_samples / config_48k.sample_rate == (
|
|
config_24k.window_samples / config_24k.sample_rate
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("algorithm", DOA_ALGORITHMS)
|
|
def test_doa_and_mvdr_use_detected_24khz_rate(algorithm: str) -> None:
|
|
config = DoaConfig(sample_rate=24_000)
|
|
estimator = DoaEstimator(config=config, algorithm=algorithm, start_worker=False)
|
|
try:
|
|
audio = synthesize_plane_wave(estimator, 68.0)
|
|
doa = estimator.estimate(audio)
|
|
assert doa.angle_degrees is not None
|
|
assert circular_error_degrees(doa.angle_degrees, 68.0) <= 1.0
|
|
|
|
mvdr = estimator._mvdr.process(audio, doa.angle_degrees)
|
|
assert estimator._mvdr.config.sample_rate == 24_000
|
|
assert mvdr.audio.shape == (config.window_samples,)
|
|
assert np.isfinite(mvdr.audio).all()
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_audio_ring_is_strictly_bounded() -> None:
|
|
config = DoaConfig(nfft=512, window_samples=1_024)
|
|
estimator = DoaEstimator(config=config, start_worker=False)
|
|
try:
|
|
oversized = np.arange(3_000 * config.channels, dtype=np.float32).reshape(
|
|
3_000, config.channels
|
|
)
|
|
estimator.submit(oversized)
|
|
assert estimator.buffered_samples == config.window_samples
|
|
np.testing.assert_array_equal(estimator._ring, oversized[-1_024:])
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_worker_publishes_doa_steered_mvdr_audio_with_one_shared_stft() -> None:
|
|
config = DoaConfig(update_interval_s=0.01, smoothing=1.0)
|
|
estimator = DoaEstimator(config=config)
|
|
try:
|
|
stft_calls = 0
|
|
analyze_stft = estimator.analyze_stft
|
|
|
|
def counted_analyze_stft(audio):
|
|
nonlocal stft_calls
|
|
stft_calls += 1
|
|
return analyze_stft(audio)
|
|
|
|
estimator.analyze_stft = counted_analyze_stft
|
|
estimator.submit(synthesize_plane_wave(estimator, 68.0))
|
|
deadline = time.monotonic() + 2.0
|
|
while estimator.latest_mvdr is None and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
|
|
doa = estimator.latest
|
|
mvdr = estimator.latest_mvdr
|
|
assert doa is not None
|
|
assert doa.angle_degrees is not None
|
|
assert circular_error_degrees(doa.angle_degrees, 68.0) <= 1.0
|
|
assert mvdr is not None
|
|
assert mvdr.audio.shape == (config.window_samples,)
|
|
assert circular_error_degrees(mvdr.angle_degrees, 68.0) <= 1.0
|
|
assert mvdr.source_sample_end == config.window_samples
|
|
queued = estimator.pop_monitor_result()
|
|
assert queued is mvdr
|
|
assert queued.stream_audio is not None
|
|
assert queued.stream_audio.size == config.window_samples - 2 * config.hop
|
|
assert stft_calls == 1
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_low_signal_keeps_mvdr_streaming_at_last_valid_direction() -> None:
|
|
config = DoaConfig(update_interval_s=0.0, smoothing=1.0)
|
|
estimator = DoaEstimator(config=config)
|
|
try:
|
|
estimator.submit(synthesize_plane_wave(estimator, 68.0))
|
|
deadline = time.monotonic() + 2.0
|
|
while estimator.latest_mvdr is None and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
|
|
first = estimator.latest_mvdr
|
|
assert first is not None
|
|
first_timestamp = first.timestamp
|
|
held_angle = first.angle_degrees
|
|
|
|
estimator.submit(
|
|
np.zeros((config.window_samples, config.channels), dtype=np.float32)
|
|
)
|
|
deadline = time.monotonic() + 2.0
|
|
while (
|
|
estimator.latest_mvdr is None
|
|
or estimator.latest_mvdr.timestamp == first_timestamp
|
|
) and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
|
|
doa = estimator.latest
|
|
held = estimator.latest_mvdr
|
|
assert doa is not None
|
|
assert doa.angle_degrees is None
|
|
assert doa.level_dbfs < config.activity_threshold_dbfs
|
|
assert held is not None
|
|
assert held.timestamp != first_timestamp
|
|
assert circular_error_degrees(held.angle_degrees, held_angle) < 1e-6
|
|
assert held.stream_audio is not None
|
|
assert held.stream_audio.size > 0
|
|
finally:
|
|
estimator.close()
|
|
|
|
|
|
def test_manual_steering_runs_mvdr_without_a_valid_doa() -> None:
|
|
config = DoaConfig(update_interval_s=0.0, smoothing=1.0)
|
|
estimator = DoaEstimator(config=config)
|
|
try:
|
|
estimator.set_mvdr_steering_angle(123.0)
|
|
estimator.submit(
|
|
np.zeros((config.window_samples, config.channels), dtype=np.float32)
|
|
)
|
|
deadline = time.monotonic() + 2.0
|
|
while estimator.latest_mvdr is None and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
|
|
doa = estimator.latest
|
|
mvdr = estimator.latest_mvdr
|
|
assert doa is not None
|
|
assert doa.angle_degrees is None
|
|
assert mvdr is not None
|
|
assert circular_error_degrees(mvdr.angle_degrees, 123.0) < 1e-6
|
|
assert mvdr.stream_audio is not None
|
|
assert mvdr.stream_audio.size > 0
|
|
finally:
|
|
estimator.close()
|