diff --git a/.env.example b/.env.example index fe55965..654625d 100644 --- a/.env.example +++ b/.env.example @@ -36,6 +36,8 @@ ENHANCE_HOP_S=0.2 BEAMFORM_MODE=auto # One BLAS/Torch thread per worker (20 workers vs 28 cores) TORCH_NUM_THREADS=1 +# Recompute GCC-PHAT TDOA every N hops (5 * 200 ms = 1 s) +BEAMFORM_TDOA_EVERY=5 # Misc PUBLISH_TIMEOUT_S=60 diff --git a/audio_cleanup.py b/audio_cleanup.py index 9d5ed88..95d43d3 100644 --- a/audio_cleanup.py +++ b/audio_cleanup.py @@ -64,7 +64,8 @@ class AudioCleaner: chunk_s: float = 1.0, hop_s: float = 0.2, beamform: str = "auto", - torch_num_threads: int = 1): + torch_num_threads: int = 1, + tdoa_every: int = 5): self.mode = mode self.beamform_mode = beamform self.sample_rate = sample_rate @@ -78,6 +79,14 @@ class AudioCleaner: self.backend = "off" self.last_dt_ms = -1.0 self.torch_num_threads = max(1, int(torch_num_threads)) + self.tdoa_every = max(1, int(tdoa_every)) + self._stft = None + self._cov = None + self._gccphat = None + self._delaysum = None + self._istft = None + self._last_tdoas = None + self._tdoa_age = 0 self._load_torch() if beamform == "force": @@ -130,12 +139,23 @@ class AudioCleaner: return try: log.info("loading speechbrain DelaySum_Beamformer (GCC-PHAT)") - self._normalizer = AudioNormalizer( - sample_rate=self.sample_rate, mix="keep") + # Native 16 kHz C920s: AudioNormalizer would resample 16k->16k. + if self.sample_rate != SAMPLE_RATE: + self._normalizer = AudioNormalizer( + sample_rate=self.sample_rate, mix="keep") + else: + self._normalizer = None bf = DelaySum_Beamformer(sampling_rate=self.sample_rate) bf.eval() self._beamformer = bf - log.info("DelaySum beamformer ready (stereo in, mono out)") + self._stft = bf.stft + self._cov = bf.cov + self._gccphat = bf.gccphat + self._delaysum = bf.delaysum + self._istft = bf.istft + log.info( + "DelaySum beamformer ready (stereo in, mono out, tdoa_every=%d)", + self.tdoa_every) except Exception as e: # noqa: BLE001 log.warning("failed to load DelaySum_Beamformer: %s", e) self._beamformer = None @@ -216,9 +236,18 @@ class AudioCleaner: 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] + # DelaySum expects [batch, time, channels] + xs = wav.unsqueeze(0) with torch.no_grad(): - y = self._beamformer(wav.unsqueeze(0)) + Xs = self._stft(xs) + if (self._last_tdoas is None + or self._tdoa_age >= self.tdoa_every): + XXs = self._cov(Xs) + self._last_tdoas = self._gccphat(XXs) + self._tdoa_age = 0 + self._tdoa_age += 1 + Ys = self._delaysum(Xs, self._last_tdoas) + y = self._istft(Ys) out = y.squeeze().detach().cpu().numpy().astype(np.float32, copy=False) if out.ndim > 1: out = np.reshape(out, -1) diff --git a/config.py b/config.py index 4e9c345..b6ca56e 100644 --- a/config.py +++ b/config.py @@ -41,6 +41,7 @@ class Config: # SpeechBrain Delay-and-Sum (GCC-PHAT) on C920 stereo before enhance beamform_mode: str = "auto" # auto | force | off torch_num_threads: int = 1 + beamform_tdoa_every: int = 5 # misc publish_timeout_s: float = 60.0 log_level: str = "INFO" @@ -88,6 +89,7 @@ def load_config(env_path: Optional[Path] = None) -> Config: enhance_hop_s=float(os.environ.get("ENHANCE_HOP_S", "0.2")), beamform_mode=os.environ.get("BEAMFORM_MODE", "auto").lower(), torch_num_threads=int(os.environ.get("TORCH_NUM_THREADS", "1")), + beamform_tdoa_every=int(os.environ.get("BEAMFORM_TDOA_EVERY", "5")), 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")), @@ -100,4 +102,6 @@ def load_config(env_path: Optional[Path] = None) -> Config: f"BEAMFORM_MODE must be auto|force|off, got {cfg.beamform_mode}") if cfg.torch_num_threads < 1: raise ValueError("TORCH_NUM_THREADS must be >= 1") + if cfg.beamform_tdoa_every < 1: + raise ValueError("BEAMFORM_TDOA_EVERY must be >= 1") return cfg diff --git a/run.py b/run.py index e1011a8..3acc038 100644 --- a/run.py +++ b/run.py @@ -44,6 +44,7 @@ def worker_args_for(cfg, cam) -> list[str]: "--enhance-hop-s", str(cfg.enhance_hop_s), "--beamform-mode", cfg.beamform_mode, "--torch-num-threads", str(cfg.torch_num_threads), + "--tdoa-every", str(cfg.beamform_tdoa_every), ] diff --git a/run_worker.py b/run_worker.py index 173c425..b30baf9 100644 --- a/run_worker.py +++ b/run_worker.py @@ -39,6 +39,7 @@ def main() -> int: ap.add_argument("--enhance-hop-s", type=float, default=0.2) ap.add_argument("--beamform-mode", default="auto") ap.add_argument("--torch-num-threads", type=int, default=1) + ap.add_argument("--tdoa-every", type=int, default=5) args = ap.parse_args() cam = json.loads(args.camera_json) @@ -65,6 +66,7 @@ def main() -> int: hop_s=args.enhance_hop_s, beamform=args.beamform_mode, torch_num_threads=args.torch_num_threads, + tdoa_every=args.tdoa_every, ) log.info("audio backend: %s", cleaner.backend) diff --git a/tests/test_audio_cleanup.py b/tests/test_audio_cleanup.py index 2dc0853..c377381 100644 --- a/tests/test_audio_cleanup.py +++ b/tests/test_audio_cleanup.py @@ -57,6 +57,32 @@ def test_beamform_loads_delaysum_module(): assert "beamform" in c.backend +def test_beamform_skips_normalizer_at_native_rate(): + c = AudioCleaner(mode="off", beamform="force", sample_rate=16000) + assert c._normalizer is None + + +def test_tdoa_reused_across_hops(): + c = AudioCleaner(mode="off", beamform="force", tdoa_every=5) + assert c._gccphat is not None + calls = {"n": 0} + orig = c._gccphat + + def wrapped(*a, **k): + calls["n"] += 1 + return orig(*a, **k) + + c._gccphat = wrapped + hop = 3200 + t = np.arange(hop) + ch0 = (np.sin(2 * np.pi * 440 * t / 16000) * 8000).astype(np.int16) + stereo = np.stack([ch0, np.roll(ch0, 2)], axis=1) + for _ in range(3): + out = c.process(stereo) + assert out.shape == (hop,) + assert calls["n"] == 1 + + def test_process_records_last_dt_ms(): c = AudioCleaner(mode="off", beamform="off") hop = 3200