143 lines
5.3 KiB
Python
143 lines
5.3 KiB
Python
"""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
|
|
|
|
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 enhancer is created once per worker process and reused across frames.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
|
|
log = logging.getLogger("audio.cleanup")
|
|
|
|
SAMPLE_RATE = 16000
|
|
|
|
|
|
class AudioCleaner:
|
|
"""Stateful audio cleaner. Call .process(int16 mono) -> int16 mono."""
|
|
|
|
def __init__(self, mode: str = "auto",
|
|
model_source: str = "speechbrain/metricgan-plus-voicebank",
|
|
sample_rate: int = SAMPLE_RATE):
|
|
self.mode = mode
|
|
self.sample_rate = sample_rate
|
|
self._enhancer = None
|
|
self._torch = None
|
|
self.backend = "off"
|
|
|
|
if mode == "force":
|
|
self._load_model(model_source)
|
|
if self.backend == "off":
|
|
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:
|
|
try:
|
|
import torch
|
|
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)
|
|
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 "
|
|
"(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)
|
|
|
|
# -- public API ---------------------------------------------------------
|
|
|
|
def process(self, pcm: np.ndarray) -> np.ndarray:
|
|
"""pcm: int16 mono numpy array. Returns int16 mono, same length.
|
|
|
|
If the model changes the length, we trim/pad back to the original
|
|
length so downstream chunking stays aligned.
|
|
"""
|
|
if self._enhancer is None:
|
|
if len(pcm) == 0:
|
|
return pcm
|
|
x = pcm.astype(np.float32) / 32768.0
|
|
y = self._simple_clean(x)
|
|
return (y * 32767.0).astype(np.float32).astype(np.int16)
|
|
|
|
x = pcm.astype(np.float32) / 32768.0
|
|
torch = self._torch
|
|
try:
|
|
w = torch.from_numpy(x).unsqueeze(0).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)
|
|
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)
|
|
out = out.squeeze(0).cpu().numpy()
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning("enhance_batch failed (%s); using fallback", e)
|
|
y = self._simple_clean(x)
|
|
return (y * 32767.0).astype(np.float32).astype(np.int16)
|
|
|
|
# align length
|
|
n = len(x)
|
|
if len(out) >= n:
|
|
out = out[:n]
|
|
else:
|
|
out = np.pad(out, (0, n - len(out)))
|
|
return (out * 32767.0).astype(np.float32).astype(np.int16)
|