From 950128deec41e976bcf5269fc81023ca7c7ca714 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 28 Aug 2026 06:16:25 +0000 Subject: [PATCH] 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. --- .env.example | 4 + README.md | 28 +++-- audio_cleanup.py | 207 +++++++++++++++++++++++++++++------- config.py | 6 ++ run.py | 1 + run_worker.py | 22 ++-- tests/test_audio_cleanup.py | 34 +++++- tests/test_config.py | 6 +- 8 files changed, 252 insertions(+), 56 deletions(-) diff --git a/.env.example b/.env.example index bb45862..f27112b 100644 --- a/.env.example +++ b/.env.example @@ -30,6 +30,10 @@ ENHANCE_MODEL=speechbrain/metricgan-plus-voicebank # model window vs publish hop: 1 s context, 200 ms emit ENHANCE_CHUNK_S=1.0 ENHANCE_HOP_S=0.2 +# C920 stereo Delay-and-Sum (SpeechBrain GCC-PHAT) before MetricGAN +# auto (default): beamform if speechbrain loads, else average channels +# force: require DelaySum_Beamformer, off: average L/R +BEAMFORM_MODE=auto # Misc PUBLISH_TIMEOUT_S=60 diff --git a/README.md b/README.md index 200269c..cfa0fee 100644 --- a/README.md +++ b/README.md @@ -71,13 +71,22 @@ re-run `--write` unless you intend to remap slots. ## Audio cleanup -The mic audio is cleaned per hop (default 200 ms) by -`speechbrain/metricgan-plus-voicebank` (16 kHz, CPU) over a 1 s rolling -window (`ENHANCE_CHUNK_S=1.0`, `ENHANCE_HOP_S=0.2`). One worker process -per camera loads the model once and reuses it. `enhance_batch` must be -called with `lengths=torch.tensor([1.0])` (relative, not sample count). -`ENHANCE_MODE=off` (or a failed model load with `auto`) falls back to a -lightweight high-pass + soft-clip DSP path so streams still work. +C920 mics are 16 kHz **stereo**. Each hop (default 200 ms) is kept +stereo through a SpeechBrain preprocessor, then published as mono: + +1. `AudioNormalizer(mix="keep")` at 16 kHz +2. Delay-and-Sum beamform (GCC-PHAT TDOA) — the live path from the + SpeechBrain [multi-microphone beamforming tutorial](https://speechbrain.readthedocs.io/en/latest/tutorials/preprocessing/multi-microphone-beamforming.html) + (`speechbrain.lobes.beamform_multimic.DelaySum_Beamformer`) +3. `speechbrain/metricgan-plus-voicebank` neural enhancement on a 1 s + rolling window (`ENHANCE_CHUNK_S=1.0`, `ENHANCE_HOP_S=0.2`) + +One worker process per camera loads the modules once and reuses them. +`enhance_batch` must be called with `lengths=torch.tensor([1.0])` +(relative, not sample count). `ENHANCE_MODE=off` (or a failed model +load with `auto`) skips MetricGAN; `BEAMFORM_MODE=off` averages L/R +instead of Delay-and-Sum. DSP high-pass + soft-clip is the last-resort +fallback so streams still work. Until the 1 s window is full, hops use the DSP fallback so LiveKit is not silent on startup. @@ -116,8 +125,9 @@ systemctl restart livekit-cameras # pick up newly plugged cameras - Video uses MJPG directly from the UVC interface (`ffmpeg -f v4l2 -input_format mjpeg`) — cheap, and converts to BGRA for LiveKit. -- Audio: ALSA mics on these C920s deliver 16 kHz **stereo** only; the worker - downmixes to mono for LiveKit. +- Audio: ALSA mics on these C920s deliver 16 kHz **stereo** only; the + worker beamforms (Delay-and-Sum) then enhances, and publishes **mono** + to LiveKit. - Camera identity is USB serial via `slots.json` / udev; USB port is fallback only for cameras not yet in the map. - Workers auto-restart (up to 10 times) if a camera drops. diff --git a/audio_cleanup.py b/audio_cleanup.py index e24e4c4..40bdeaa 100644 --- a/audio_cleanup.py +++ b/audio_cleanup.py @@ -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) diff --git a/config.py b/config.py index c982541..7047615 100644 --- a/config.py +++ b/config.py @@ -38,6 +38,8 @@ class Config: enhance_model: str = "speechbrain/metricgan-plus-voicebank" enhance_chunk_s: float = 1.0 enhance_hop_s: float = 0.2 + # SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance + beamform_mode: str = "auto" # auto | force | off # misc publish_timeout_s: float = 60.0 log_level: str = "INFO" @@ -83,6 +85,7 @@ def load_config(env_path: Optional[Path] = None) -> Config: ), enhance_chunk_s=float(os.environ.get("ENHANCE_CHUNK_S", "1.0")), enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")), + beamform_mode=os.environ.get("BEAMFORM_MODE", "auto").lower(), publish_timeout_s=float(os.environ.get("PUBLISH_TIMEOUT_S", "60")), log_level=os.environ.get("LOG_LEVEL", "INFO").upper(), cameras=_parse_cameras(os.environ.get("CAMERAS", "all")), @@ -90,4 +93,7 @@ def load_config(env_path: Optional[Path] = None) -> Config: if cfg.enhance_mode not in ("auto", "force", "off"): raise ValueError(f"ENHANCE_MODE must be auto|force|off, got {cfg.enhance_mode}") + if cfg.beamform_mode not in ("auto", "force", "off"): + raise ValueError( + f"BEAMFORM_MODE must be auto|force|off, got {cfg.beamform_mode}") return cfg diff --git a/run.py b/run.py index ac73312..b3c6fed 100644 --- a/run.py +++ b/run.py @@ -42,6 +42,7 @@ def worker_args_for(cfg, cam) -> list[str]: "--enhance-model", cfg.enhance_model, "--enhance-chunk-s", str(cfg.enhance_chunk_s), "--enhance-hop-s", str(cfg.enhance_hop_s), + "--beamform-mode", cfg.beamform_mode, ] diff --git a/run_worker.py b/run_worker.py index 3d7f67f..a832f0d 100644 --- a/run_worker.py +++ b/run_worker.py @@ -37,6 +37,7 @@ def main() -> int: default="speechbrain/metricgan-plus-voicebank") ap.add_argument("--enhance-chunk-s", type=float, default=1.0) ap.add_argument("--enhance-hop-s", type=float, default=0.2) + ap.add_argument("--beamform-mode", default="auto") args = ap.parse_args() cam = json.loads(args.camera_json) @@ -60,6 +61,7 @@ def main() -> int: sample_rate=args.audio_rate, chunk_s=args.enhance_chunk_s, hop_s=args.enhance_hop_s, + beamform=args.beamform_mode, ) log.info("audio backend: %s", cleaner.backend) @@ -138,6 +140,7 @@ def main() -> int: # chunk * 2 ch * 2 bytes need = chunk * 2 * 2 buf = b"" + hops = 0 while aproc.poll() is None: data = await loop.run_in_executor( None, aproc.stdout.read, 4096) @@ -146,21 +149,28 @@ def main() -> int: buf += data while len(buf) >= need: piece, buf = buf[:need], buf[need:] - pcm_stereo = np.frombuffer(piece, dtype=np.int16) - pcm_mono = pcm_stereo.reshape(-1, 2).mean( - axis=1).astype(np.int16) - # align to full chunks - usable = (len(pcm_mono) // chunk) * chunk + pcm_stereo = np.frombuffer( + piece, dtype=np.int16).reshape(-1, 2) + n_frames = pcm_stereo.shape[0] + usable = (n_frames // chunk) * chunk if usable: cleaned = await loop.run_in_executor( None, cleaner.process, - pcm_mono[:usable]) + pcm_stereo[:usable]) + t_cap = time.perf_counter() await audio_source.capture_frame( rtc.AudioFrame( cleaned.tobytes(), args.audio_rate, 1, len(cleaned))) + hops += 1 + if hops % 50 == 0: + log.info( + "audio hop process_ms=%.1f capture_ms=%.1f", + cleaner.last_dt_ms, + (time.perf_counter() - t_cap) * 1000.0, + ) aproc = subprocess.Popen( ["ffmpeg", "-hide_banner", "-loglevel", "error", diff --git a/tests/test_audio_cleanup.py b/tests/test_audio_cleanup.py index a8643b9..2dc0853 100644 --- a/tests/test_audio_cleanup.py +++ b/tests/test_audio_cleanup.py @@ -8,7 +8,7 @@ from audio_cleanup import AudioCleaner def test_enhance_batch_without_lengths_typeerror(): c = AudioCleaner(mode="force") - assert c.backend == "model" + assert "model" in c.backend import torch w = torch.zeros(1, 8000) @@ -17,7 +17,7 @@ def test_enhance_batch_without_lengths_typeerror(): def test_process_uses_relative_lengths_one(): - src = inspect.getsource(AudioCleaner.process) + src = inspect.getsource(AudioCleaner._process_inner) assert "tensor([1.0])" in src @@ -48,3 +48,33 @@ def test_overlap_runs_model_after_enough_context(): out = c.process(pcm) assert out.shape == (hop,) assert out.dtype == np.int16 + + +def test_beamform_loads_delaysum_module(): + c = AudioCleaner(mode="off", beamform="force") + assert c._beamformer is not None + assert type(c._beamformer).__name__ == "DelaySum_Beamformer" + assert "beamform" in c.backend + + +def test_process_records_last_dt_ms(): + c = AudioCleaner(mode="off", beamform="off") + hop = 3200 + pcm = (np.random.randn(hop) * 500).astype(np.int16) + c.process(pcm) + assert c.last_dt_ms >= 0.0 + + +def test_stereo_beamform_then_model_preserves_hop(): + c = AudioCleaner( + mode="force", beamform="force", chunk_s=0.5, hop_s=0.25) + assert c.backend == "model+beamform" + hop = 4000 + t = np.arange(hop) + ch0 = (np.sin(2 * np.pi * 440 * t / 16000) * 8000).astype(np.int16) + ch1 = np.roll(ch0, 3) + stereo = np.stack([ch0, ch1], axis=1) + for _ in range(4): + out = c.process(stereo) + assert out.dtype == np.int16 + assert out.shape == (hop,) diff --git a/tests/test_config.py b/tests/test_config.py index bff5236..5c05441 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,7 +1,11 @@ -from config import _parse_cameras +from config import Config, _parse_cameras def test_parse_cameras_list(): assert _parse_cameras("1,5") == ["1", "5"] assert _parse_cameras("all") == ["all"] assert _parse_cameras("1-3") == ["1", "2", "3"] + + +def test_default_beamform_mode_auto(): + assert Config().beamform_mode == "auto"