feat: DelaySum beamform plus hop process latency
Keep C920 stereo through SpeechBrain GCC-PHAT Delay-and-Sum, then MetricGAN, and publish mono to LiveKit. Record last_dt_ms and log process_ms/capture_ms every 50 hops for the latency plan.
This commit is contained in:
root
+169
-38
@@ -1,23 +1,33 @@
|
||||
"""Audio cleanup for camera mics, using speechbrain.
|
||||
|
||||
Pipeline per audio frame (16 kHz mono int16, ~100 ms chunks from the mic):
|
||||
int16 -> float32 [-1, 1] -> speechbrain enhancement -> int16
|
||||
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
|
||||
|
||||
Two implementations, selected by availability and the ENHANCE_MODE config:
|
||||
* "model" - speechbrain SpectralMaskEnhancement (metricGAN-plus). Real
|
||||
ML-based noise suppression. Runs on CPU; ~20-40 ms per 100 ms
|
||||
chunk on a modern core, so latency stays well under a frame
|
||||
budget if workers are spread across cores.
|
||||
* "gain" - lightweight DSP fallback (soft-clip + high-pass) used when the
|
||||
model can't be loaded (ENHANCE_MODE=off / auto + no model).
|
||||
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).
|
||||
|
||||
The enhancer is created once per worker process and reused across frames.
|
||||
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
|
||||
from typing import Optional
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -27,35 +37,98 @@ SAMPLE_RATE = 16000
|
||||
|
||||
|
||||
class AudioCleaner:
|
||||
"""Stateful audio cleaner. Call .process(int16 mono) -> int16 mono."""
|
||||
"""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):
|
||||
hop_s: float = 0.2,
|
||||
beamform: str = "auto"):
|
||||
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._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.backend == "off":
|
||||
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)
|
||||
# mode == "off": passthrough + basic gain shaping
|
||||
|
||||
def _load_model(self, source: str) -> None:
|
||||
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:
|
||||
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,
|
||||
)
|
||||
@@ -65,14 +138,12 @@ class AudioCleaner:
|
||||
return
|
||||
try:
|
||||
log.info("loading speechbrain enhancement model %s ...", source)
|
||||
self._torch = torch
|
||||
model = SpectralMaskEnhancement.from_hparams(
|
||||
source=source,
|
||||
savedir="/root/.cache/speechbrain-enhancement",
|
||||
)
|
||||
model.mods.enhance_model.eval()
|
||||
self._enhancer = model
|
||||
self.backend = "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 "
|
||||
@@ -102,52 +173,112 @@ class AudioCleaner:
|
||||
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 numpy array. Returns int16 mono, same length.
|
||||
"""pcm: int16 mono [N] or stereo [N, 2]. Returns int16 mono [N].
|
||||
|
||||
Accumulates a `chunk_s` window and runs the 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.
|
||||
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.
|
||||
"""
|
||||
n = len(pcm)
|
||||
if n == 0:
|
||||
return pcm
|
||||
x = pcm.astype(np.float32) / 32768.0
|
||||
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:
|
||||
y = self._simple_clean(x)
|
||||
return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
|
||||
return self._to_int16(self._simple_clean(mono_hop))
|
||||
|
||||
self._ctx = np.concatenate([self._ctx, x])
|
||||
self._ctx = np.concatenate([self._ctx, mono_hop])
|
||||
if len(self._ctx) < self.chunk:
|
||||
y = self._simple_clean(x)
|
||||
return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
|
||||
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).to(
|
||||
self._enhancer.device)
|
||||
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(self._enhancer.device)
|
||||
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).cpu().numpy()
|
||||
enhanced = out.squeeze(0).detach().cpu().numpy()
|
||||
except Exception as e: # noqa: BLE001
|
||||
log.warning("enhance_batch failed (%s); using fallback", e)
|
||||
y = self._simple_clean(x)
|
||||
return (np.clip(y, -1.0, 1.0) * 32767.0).astype(np.int16)
|
||||
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 (np.clip(hop_out, -1.0, 1.0) * 32767.0).astype(np.int16)
|
||||
return self._to_int16(hop_out)
|
||||
|
||||
Reference in New Issue
Block a user