162 lines
5.2 KiB
Python
162 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from src.doa_estimator import MIC_AZIMUTHS_DEG
|
|
from src.mvdr_beamformer import MvdrBeamformer, MvdrConfig
|
|
from src.stft_analysis import analyze_multichannel_stft
|
|
|
|
|
|
SAMPLE_RATE = 48_000
|
|
SAMPLES = 8_192
|
|
|
|
|
|
def microphone_positions() -> np.ndarray:
|
|
angles = np.deg2rad(np.asarray(MIC_AZIMUTHS_DEG))
|
|
return 0.056 * np.vstack((np.cos(angles), np.sin(angles)))
|
|
|
|
|
|
def band_limited_noise(rng: np.random.Generator) -> np.ndarray:
|
|
source = rng.standard_normal(SAMPLES)
|
|
frequencies = np.fft.rfftfreq(SAMPLES, 1.0 / SAMPLE_RATE)
|
|
spectrum = np.fft.rfft(source)
|
|
spectrum[(frequencies < 700.0) | (frequencies > 3_000.0)] = 0.0
|
|
filtered = np.fft.irfft(spectrum, n=SAMPLES)
|
|
return filtered / np.std(filtered)
|
|
|
|
|
|
def plane_wave(source: np.ndarray, angle_degrees: float) -> np.ndarray:
|
|
positions = microphone_positions()
|
|
frequencies = np.fft.rfftfreq(source.size, 1.0 / SAMPLE_RATE)
|
|
angle = np.deg2rad(angle_degrees)
|
|
direction = np.array([np.cos(angle), np.sin(angle)])
|
|
relative_times = direction @ positions / 343.0
|
|
spectrum = np.fft.rfft(source)
|
|
return np.stack(
|
|
[
|
|
np.fft.irfft(
|
|
spectrum * np.exp(2j * np.pi * frequencies * delay),
|
|
n=source.size,
|
|
)
|
|
for delay in relative_times
|
|
],
|
|
axis=1,
|
|
)
|
|
|
|
|
|
def test_mvdr_preserves_signal_from_steered_direction() -> None:
|
|
rng = np.random.default_rng(21)
|
|
source = band_limited_noise(rng)
|
|
audio = 0.05 * plane_wave(source, 73.0)
|
|
beamformer = MvdrBeamformer(microphone_positions(), MvdrConfig())
|
|
|
|
result = beamformer.process(audio, 73.0)
|
|
|
|
valid = slice(512, -512)
|
|
gain = np.linalg.lstsq(
|
|
source[valid, None], result.audio[valid], rcond=None
|
|
)[0][0]
|
|
assert np.isfinite(result.audio).all()
|
|
assert result.audio.shape == (SAMPLES,)
|
|
assert not result.audio.flags.writeable
|
|
assert result.angle_degrees == 73.0
|
|
assert 0.045 < gain < 0.055
|
|
|
|
|
|
def test_mvdr_suppresses_off_axis_interference() -> None:
|
|
rng = np.random.default_rng(17)
|
|
target = band_limited_noise(rng)
|
|
interference = band_limited_noise(rng)
|
|
mixture = 0.05 * (
|
|
plane_wave(target, 40.0) + plane_wave(interference, 145.0)
|
|
)
|
|
beamformer = MvdrBeamformer(microphone_positions(), MvdrConfig())
|
|
|
|
output = beamformer.process(mixture, 40.0).audio
|
|
|
|
valid = slice(512, -512)
|
|
references = np.column_stack((target[valid], interference[valid]))
|
|
input_coefficients = np.linalg.lstsq(
|
|
references, mixture[valid, 0], rcond=None
|
|
)[0]
|
|
output_coefficients = np.linalg.lstsq(
|
|
references, output[valid], rcond=None
|
|
)[0]
|
|
input_sir = 20.0 * np.log10(
|
|
abs(input_coefficients[0] / input_coefficients[1])
|
|
)
|
|
output_sir = 20.0 * np.log10(
|
|
abs(output_coefficients[0] / output_coefficients[1])
|
|
)
|
|
|
|
assert output_sir >= input_sir + 10.0
|
|
|
|
|
|
def test_streaming_wola_excludes_unstable_edges_and_stays_continuous() -> None:
|
|
sample_rate = 24_000
|
|
nfft = 256
|
|
hop = 128
|
|
window_samples = 4_096
|
|
step = 3_072
|
|
total = window_samples + 4 * step
|
|
rng = np.random.default_rng(123)
|
|
frequencies = np.fft.rfftfreq(total, 1.0 / sample_rate)
|
|
passband = (frequencies >= 600.0) & (frequencies <= 3_800.0)
|
|
|
|
def noise() -> np.ndarray:
|
|
spectrum = np.fft.rfft(rng.standard_normal(total))
|
|
spectrum[~passband] = 0.0
|
|
signal = np.fft.irfft(spectrum, total)
|
|
return signal / np.std(signal)
|
|
|
|
positions = microphone_positions()
|
|
|
|
def wave(source: np.ndarray, angle_degrees: float) -> np.ndarray:
|
|
angle = np.deg2rad(angle_degrees)
|
|
direction = np.array([np.cos(angle), np.sin(angle)])
|
|
delays = direction @ positions / 343.0
|
|
spectrum = np.fft.rfft(source)
|
|
return np.stack(
|
|
[
|
|
np.fft.irfft(
|
|
spectrum * np.exp(2j * np.pi * frequencies * delay),
|
|
total,
|
|
)
|
|
for delay in delays
|
|
],
|
|
axis=1,
|
|
)
|
|
|
|
audio = 0.04 * (wave(noise(), 40.0) + wave(noise(), 145.0))
|
|
config = MvdrConfig(sample_rate=sample_rate, nfft=nfft, hop=hop)
|
|
beamformer = MvdrBeamformer(positions, config)
|
|
window = np.hanning(nfft).astype(np.float32)
|
|
stream_blocks = []
|
|
previous_end = None
|
|
|
|
for start in range(0, total - window_samples + 1, step):
|
|
analysis = analyze_multichannel_stft(
|
|
audio[start : start + window_samples],
|
|
channels=8,
|
|
nfft=nfft,
|
|
hop=hop,
|
|
window=window,
|
|
)
|
|
result = beamformer.process_stft(
|
|
analysis,
|
|
40.0,
|
|
source_sample_end=start + window_samples,
|
|
)
|
|
assert result.stream_audio is not None
|
|
assert result.stream_audio.size > 0
|
|
if previous_end is not None:
|
|
assert result.stream_sample_start == previous_end
|
|
assert result.stream_audio.size == step
|
|
previous_end = result.stream_sample_end
|
|
stream_blocks.append(result.stream_audio)
|
|
|
|
joined = np.concatenate(stream_blocks)
|
|
assert stream_blocks[0].size == window_samples - 2 * hop
|
|
assert np.max(np.abs(joined)) < 0.5
|
|
assert np.isfinite(joined).all()
|