105 lines
4.2 KiB
Python
105 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
from scipy import signal
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ReferenceSuppressionReport:
|
|
channel: int
|
|
input_rms: float
|
|
predicted_rms: float
|
|
output_rms: float
|
|
reduction_db: float
|
|
phase_mapping: tuple[tuple[float, float, float], ...]
|
|
|
|
|
|
def suppress_program_reference(
|
|
field: np.ndarray,
|
|
reference: np.ndarray,
|
|
sample_rate: int,
|
|
channels: list[int] | tuple[int, ...],
|
|
*,
|
|
nperseg: int = 2048,
|
|
noverlap: int = 1536,
|
|
regularization: float = 1e-6,
|
|
subtraction_strength: float = 0.85,
|
|
) -> tuple[np.ndarray, list[ReferenceSuppressionReport]]:
|
|
x = np.asarray(field, dtype=np.float32)
|
|
ref = np.asarray(reference, dtype=np.float32).reshape(-1)
|
|
if x.ndim != 2:
|
|
raise ValueError("field must be frames-by-channels")
|
|
ref = pad_or_trim(ref, len(x))
|
|
cleaned = x.copy()
|
|
reports = []
|
|
for channel in channels:
|
|
ch = int(channel)
|
|
if ch < 0 or ch >= x.shape[1]:
|
|
continue
|
|
transfer, freqs = estimate_transfer(reference=ref, observed=x[:, ch], sample_rate=sample_rate, nperseg=nperseg, noverlap=noverlap, regularization=regularization)
|
|
predicted = synthesize_reference(reference=ref, transfer=transfer, sample_rate=sample_rate, nperseg=nperseg, noverlap=noverlap)
|
|
predicted = pad_or_trim(predicted, len(x))
|
|
before = x[:, ch]
|
|
after = before - float(subtraction_strength) * predicted
|
|
cleaned[:, ch] = after.astype(np.float32)
|
|
before_rms = rms(before)
|
|
after_rms = rms(after)
|
|
reports.append(
|
|
ReferenceSuppressionReport(
|
|
channel=ch,
|
|
input_rms=before_rms,
|
|
predicted_rms=rms(predicted),
|
|
output_rms=after_rms,
|
|
reduction_db=20.0 * np.log10(max(after_rms, 1e-12) / max(before_rms, 1e-12)),
|
|
phase_mapping=phase_mapping(freqs, transfer),
|
|
)
|
|
)
|
|
return cleaned.astype(np.float32), reports
|
|
|
|
|
|
def estimate_transfer(*, reference: np.ndarray, observed: np.ndarray, sample_rate: int, nperseg: int, noverlap: int, regularization: float):
|
|
freqs, _, ref_stft = signal.stft(reference, fs=sample_rate, nperseg=nperseg, noverlap=noverlap, boundary="zeros", padded=True)
|
|
_, _, obs_stft = signal.stft(pad_or_trim(observed, len(reference)), fs=sample_rate, nperseg=nperseg, noverlap=noverlap, boundary="zeros", padded=True)
|
|
numerator = np.sum(obs_stft * np.conj(ref_stft), axis=1)
|
|
denominator = np.sum(np.abs(ref_stft) ** 2, axis=1) + float(regularization)
|
|
transfer = numerator / denominator
|
|
return transfer.astype(np.complex64), freqs.astype(np.float32)
|
|
|
|
|
|
def synthesize_reference(*, reference: np.ndarray, transfer: np.ndarray, sample_rate: int, nperseg: int, noverlap: int):
|
|
_, _, ref_stft = signal.stft(reference, fs=sample_rate, nperseg=nperseg, noverlap=noverlap, boundary="zeros", padded=True)
|
|
predicted_stft = ref_stft * transfer[:, None]
|
|
_, predicted = signal.istft(predicted_stft, fs=sample_rate, nperseg=nperseg, noverlap=noverlap, input_onesided=True, boundary=True)
|
|
return predicted.astype(np.float32)
|
|
|
|
|
|
def phase_mapping(freqs: np.ndarray, transfer: np.ndarray, *, min_hz: float = 80.0, max_hz: float = 16000.0, bins: int = 24):
|
|
valid = (freqs >= min_hz) & (freqs <= max_hz) & (np.abs(transfer) > 1e-9)
|
|
if not np.any(valid):
|
|
return tuple()
|
|
selected_freqs = np.geomspace(max(min_hz, float(freqs[valid][0])), min(max_hz, float(freqs[valid][-1])), bins)
|
|
rows = []
|
|
for frequency in selected_freqs:
|
|
index = int(np.argmin(np.abs(freqs - frequency)))
|
|
magnitude = float(abs(transfer[index]))
|
|
phase = float(np.angle(transfer[index]))
|
|
rows.append((float(freqs[index]), phase, magnitude))
|
|
return tuple(rows)
|
|
|
|
|
|
def pad_or_trim(samples: np.ndarray, length: int) -> np.ndarray:
|
|
x = np.asarray(samples, dtype=np.float32).reshape(-1)
|
|
if len(x) == length:
|
|
return x
|
|
if len(x) > length:
|
|
return x[:length]
|
|
out = np.zeros(length, dtype=np.float32)
|
|
out[: len(x)] = x
|
|
return out
|
|
|
|
|
|
def rms(samples: np.ndarray) -> float:
|
|
x = np.asarray(samples, dtype=np.float32)
|
|
return float(np.sqrt(np.mean(x * x))) if x.size else 0.0
|