root
23c2e9c4c0
Default TORCH_NUM_THREADS=1 so 14-20 workers do not each spawn a 28-thread BLAS pool.
303 lines
11 KiB
Python
303 lines
11 KiB
Python
"""Audio cleanup for camera mics, using speechbrain.
|
|
|
|
Pipeline per audio hop (16 kHz from the C920; stereo ALSA, LiveKit wants mono):
|
|
int16 stereo -> float32 [-1, 1]
|
|
-> SpeechBrain AudioNormalizer (keep channels, 16 kHz)
|
|
-> Delay-and-Sum beamform (tutorial recipe: STFT + Covariance +
|
|
GCC-PHAT TDOA + DelaySum + ISTFT) -> mono
|
|
-> MetricGAN-plus neural enhancement
|
|
-> int16 mono
|
|
|
|
The Delay-and-Sum path matches
|
|
https://speechbrain.readthedocs.io/en/latest/tutorials/preprocessing/multi-microphone-beamforming.html
|
|
("Delay-and-Sum Beamforming" / GCC-PHAT). Wrapped as
|
|
speechbrain.lobes.beamform_multimic.DelaySum_Beamformer (a torch.nn.Module).
|
|
|
|
MVDR / GEV from that tutorial need a separate noise covariance, which a
|
|
live C920 hop does not have, so they are not used online.
|
|
|
|
Two enhancement implementations, selected by ENHANCE_MODE:
|
|
* "model" - speechbrain SpectralMaskEnhancement (metricGAN-plus)
|
|
* "gain" - high-pass + soft-clip when the model is unavailable
|
|
|
|
Beamforming is selected independently by BEAMFORM_MODE (auto|force|off).
|
|
The preprocessor is created once per worker process and reused.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
log = logging.getLogger("audio.cleanup")
|
|
|
|
SAMPLE_RATE = 16000
|
|
|
|
|
|
def apply_torch_thread_limits(n: int = 1) -> None:
|
|
"""Cap BLAS/Torch pools so N camera workers do not oversubscribe."""
|
|
n = max(1, int(n))
|
|
os.environ.setdefault("OMP_NUM_THREADS", str(n))
|
|
os.environ.setdefault("MKL_NUM_THREADS", str(n))
|
|
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(n))
|
|
try:
|
|
import torch
|
|
torch.set_num_threads(n)
|
|
torch.set_num_interop_threads(1)
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
|
|
|
|
class AudioCleaner:
|
|
"""Stateful audio cleaner. Call .process(...) -> int16 mono.
|
|
|
|
Accepts int16 mono ``[N]`` or stereo ``[N, 2]`` (C920 ALSA). Returns
|
|
int16 mono of length N for LiveKit.
|
|
"""
|
|
|
|
def __init__(self, mode: str = "auto",
|
|
model_source: str = "speechbrain/metricgan-plus-voicebank",
|
|
sample_rate: int = SAMPLE_RATE,
|
|
chunk_s: float = 1.0,
|
|
hop_s: float = 0.2,
|
|
beamform: str = "auto",
|
|
torch_num_threads: int = 1):
|
|
self.mode = mode
|
|
self.beamform_mode = beamform
|
|
self.sample_rate = sample_rate
|
|
self.chunk = max(1, int(chunk_s * sample_rate))
|
|
self.hop = max(1, int(hop_s * sample_rate))
|
|
self._ctx = np.zeros(0, dtype=np.float32)
|
|
self._enhancer = None
|
|
self._beamformer = None
|
|
self._normalizer = None
|
|
self._torch = None
|
|
self.backend = "off"
|
|
self.last_dt_ms = -1.0
|
|
self.torch_num_threads = max(1, int(torch_num_threads))
|
|
|
|
self._load_torch()
|
|
if beamform == "force":
|
|
self._load_beamformer()
|
|
if self._beamformer is None:
|
|
raise RuntimeError(
|
|
"BEAMFORM_MODE=force but failed to load "
|
|
"SpeechBrain DelaySum_Beamformer")
|
|
elif beamform == "auto":
|
|
self._load_beamformer()
|
|
|
|
if mode == "force":
|
|
self._load_model(model_source)
|
|
if self._enhancer is None:
|
|
raise RuntimeError(
|
|
f"ENHANCE_MODE=force but failed to load model "
|
|
f"'{model_source}'")
|
|
elif mode == "auto":
|
|
self._load_model(model_source)
|
|
|
|
self._refresh_backend()
|
|
|
|
def _refresh_backend(self) -> None:
|
|
parts = []
|
|
if self._enhancer is not None:
|
|
parts.append("model")
|
|
if self._beamformer is not None:
|
|
parts.append("beamform")
|
|
self.backend = "+".join(parts) if parts else "off"
|
|
|
|
def _load_torch(self) -> None:
|
|
apply_torch_thread_limits(self.torch_num_threads)
|
|
try:
|
|
import torch
|
|
self._torch = torch
|
|
except Exception as e: # noqa: BLE001
|
|
log.info("torch unavailable (%s)", e)
|
|
self._torch = None
|
|
|
|
def _load_beamformer(self) -> None:
|
|
"""Tutorial Delay-and-Sum: STFT, Covariance, GCC-PHAT, DelaySum, ISTFT."""
|
|
if self._torch is None:
|
|
log.info("no torch; skipping DelaySum beamformer")
|
|
return
|
|
try:
|
|
from speechbrain.dataio.preprocess import AudioNormalizer
|
|
from speechbrain.lobes.beamform_multimic import DelaySum_Beamformer
|
|
except Exception as e: # noqa: BLE001
|
|
log.info("speechbrain beamformer unavailable (%s)", e)
|
|
return
|
|
try:
|
|
log.info("loading speechbrain DelaySum_Beamformer (GCC-PHAT)")
|
|
self._normalizer = AudioNormalizer(
|
|
sample_rate=self.sample_rate, mix="keep")
|
|
bf = DelaySum_Beamformer(sampling_rate=self.sample_rate)
|
|
bf.eval()
|
|
self._beamformer = bf
|
|
log.info("DelaySum beamformer ready (stereo in, mono out)")
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning("failed to load DelaySum_Beamformer: %s", e)
|
|
self._beamformer = None
|
|
self._normalizer = None
|
|
|
|
def _load_model(self, source: str) -> None:
|
|
if self._torch is None:
|
|
log.info("torch unavailable; using fallback")
|
|
self._enhancer = None
|
|
return
|
|
try:
|
|
from speechbrain.inference.enhancement import (
|
|
SpectralMaskEnhancement,
|
|
)
|
|
except Exception as e: # noqa: BLE001
|
|
log.info("speechbrain unavailable (%s); using fallback", e)
|
|
self._enhancer = None
|
|
return
|
|
try:
|
|
log.info("loading speechbrain enhancement model %s ...", source)
|
|
model = SpectralMaskEnhancement.from_hparams(
|
|
source=source,
|
|
savedir="/root/.cache/speechbrain-enhancement",
|
|
)
|
|
model.mods.enhance_model.eval()
|
|
self._enhancer = model
|
|
log.info("speechbrain enhancement ready (16 kHz in/out)")
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning("failed to load enhancement model %s: %s "
|
|
"(using fallback)", source, e)
|
|
self._enhancer = None
|
|
|
|
# -- DSP fallback ------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _simple_clean(x: np.ndarray) -> np.ndarray:
|
|
"""Lightweight cleanup: high-pass + soft knee + gentle gain.
|
|
|
|
Good enough to remove DC offset/low rumble and taming clipping
|
|
without any ML model.
|
|
"""
|
|
# first-order high-pass at ~80 Hz
|
|
if len(x) > 1:
|
|
alpha = 0.998
|
|
y = np.empty_like(x)
|
|
y[0] = x[0]
|
|
for i in range(1, len(x)):
|
|
y[i] = alpha * (y[i - 1] + x[i] - x[i - 1])
|
|
x = y
|
|
# soft-clip to avoid harsh distortion
|
|
x = np.tanh(x * 0.9) / np.tanh(0.9)
|
|
# gentle makeup gain
|
|
x = x * 1.1
|
|
return np.clip(x, -1.0, 1.0)
|
|
|
|
# -- beamform (tutorial Delay-and-Sum) --------------------------------
|
|
|
|
def _to_mono(self, x: np.ndarray) -> np.ndarray:
|
|
"""x: float32 [T, C] -> float32 [T]."""
|
|
if x.shape[1] == 1:
|
|
return x[:, 0]
|
|
if self._beamformer is None or self._torch is None:
|
|
return x.mean(axis=1)
|
|
try:
|
|
return self._beamform(x)
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning("DelaySum beamform failed (%s); averaging channels", e)
|
|
return x.mean(axis=1)
|
|
|
|
def _beamform(self, stereo: np.ndarray) -> np.ndarray:
|
|
"""SpeechBrain Delay-and-Sum on [T, 2] float32 -> [T] float32.
|
|
|
|
Recipe from the multi-microphone beamforming tutorial:
|
|
STFT -> Covariance -> GCC-PHAT TDOA -> DelaySum -> ISTFT.
|
|
"""
|
|
torch = self._torch
|
|
n = stereo.shape[0]
|
|
wav = torch.from_numpy(np.ascontiguousarray(stereo, dtype=np.float32))
|
|
if self._normalizer is not None:
|
|
wav = self._normalizer(wav, self.sample_rate)
|
|
# DelaySum_Beamformer expects [batch, time, channels]
|
|
with torch.no_grad():
|
|
y = self._beamformer(wav.unsqueeze(0))
|
|
out = y.squeeze().detach().cpu().numpy().astype(np.float32, copy=False)
|
|
if out.ndim > 1:
|
|
out = np.reshape(out, -1)
|
|
if out.shape[0] != n:
|
|
if out.shape[0] > n:
|
|
out = out[:n]
|
|
else:
|
|
out = np.pad(out, (0, n - out.shape[0]))
|
|
return out
|
|
|
|
def _to_int16(self, y: np.ndarray) -> np.ndarray:
|
|
return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
|
|
|
|
# -- public API ---------------------------------------------------------
|
|
|
|
def process(self, pcm: np.ndarray) -> np.ndarray:
|
|
"""pcm: int16 mono [N] or stereo [N, 2]. Returns int16 mono [N].
|
|
|
|
Accumulates a `chunk_s` window and runs beamform + model on that
|
|
window, emitting only the latest hop (the input length). Until a
|
|
full window is buffered, the DSP fallback is used so the stream
|
|
is not silent.
|
|
"""
|
|
pcm = np.asarray(pcm)
|
|
t0 = time.perf_counter()
|
|
try:
|
|
return self._process_inner(pcm)
|
|
finally:
|
|
self.last_dt_ms = (time.perf_counter() - t0) * 1000.0
|
|
|
|
def _process_inner(self, pcm: np.ndarray) -> np.ndarray:
|
|
if pcm.size == 0:
|
|
return np.zeros(0, dtype=np.int16)
|
|
if pcm.ndim == 1:
|
|
frames = pcm.reshape(-1, 1)
|
|
elif pcm.ndim == 2:
|
|
frames = pcm
|
|
else:
|
|
raise ValueError(f"pcm must be 1-D or 2-D, got shape {pcm.shape}")
|
|
|
|
n = frames.shape[0]
|
|
x = frames.astype(np.float32) / 32768.0
|
|
|
|
# Delay-and-Sum the *hop* (200 ms is ~60 ms here). Beamforming the
|
|
# full 1 s window every hop is ~330 ms — not 20-camera real-time.
|
|
# Same tutorial modules; TDOA is estimated per hop via GCC-PHAT.
|
|
mono_hop = self._to_mono(x)
|
|
|
|
if self._enhancer is None:
|
|
return self._to_int16(self._simple_clean(mono_hop))
|
|
|
|
self._ctx = np.concatenate([self._ctx, mono_hop])
|
|
if len(self._ctx) < self.chunk:
|
|
return self._to_int16(self._simple_clean(mono_hop))
|
|
self._ctx = self._ctx[-self.chunk:]
|
|
|
|
torch = self._torch
|
|
try:
|
|
w = torch.from_numpy(self._ctx.copy()).unsqueeze(0)
|
|
if hasattr(self._enhancer, "device"):
|
|
w = w.to(self._enhancer.device)
|
|
# speechbrain >= 1.1 (e.g. metricGAN-plus) requires `lengths`
|
|
# on enhance_batch. The value is RELATIVE (multiplied by the
|
|
# number of frames internally), so 1.0 = "the whole sequence".
|
|
# Older versions accept the call without it (TypeError fallback).
|
|
lengths = torch.tensor([1.0]).to(w.device)
|
|
try:
|
|
with torch.no_grad():
|
|
out = self._enhancer.enhance_batch(w, lengths=lengths)
|
|
except TypeError:
|
|
with torch.no_grad():
|
|
out = self._enhancer.enhance_batch(w)
|
|
enhanced = out.squeeze(0).detach().cpu().numpy()
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning("enhance_batch failed (%s); using fallback", e)
|
|
return self._to_int16(self._simple_clean(mono_hop))
|
|
|
|
hop_out = enhanced[-n:] if len(enhanced) >= n else np.pad(
|
|
enhanced, (0, n - len(enhanced)))
|
|
hop_out = hop_out[:n]
|
|
return self._to_int16(hop_out)
|