fix(onnx): handle short chunks and short streams in Fsmn_vad_online

.github/workflows/test-funasr-onnx-package.yml

@@ -11,6 +11,7 @@ on:
       - "tests/test_funasr_onnx_installed.py"
       - "tests/test_funasr_onnx_paraformer_online_overlap.py"
       - "tests/test_funasr_onnx_paraformer_online_empty_final.py"
+      - "tests/test_funasr_onnx_vad_online_short_chunks.py"
       - "tests/test_rich_transcription_postprocess.py"
       - "tests/test_rich_transcription_event_segments.py"
   push:
@@ -25,6 +26,7 @@ on:
       - "tests/test_funasr_onnx_installed.py"
       - "tests/test_funasr_onnx_paraformer_online_overlap.py"
       - "tests/test_funasr_onnx_paraformer_online_empty_final.py"
+      - "tests/test_funasr_onnx_vad_online_short_chunks.py"
       - "tests/test_rich_transcription_postprocess.py"
       - "tests/test_rich_transcription_event_segments.py"
   workflow_dispatch:
@@ -105,6 +107,9 @@ jobs:
       - name: Verify empty split-final Paraformer result
         run: python tests/test_funasr_onnx_paraformer_online_empty_final.py
 
+      - name: Verify streaming VAD short chunks
+        run: python tests/test_funasr_onnx_vad_online_short_chunks.py
+
       - name: Verify source distribution error paths
         run: |
           python -m pip install --force-reinstall --no-deps dist/funasr-onnx/*.tar.gz

runtime/python/onnxruntime/funasr_onnx/utils/e2e_vad.py

@@ -452,6 +452,10 @@ def LatencyFrmNumAtStartPoint(self) -> int:
 
     def GetFrameState(self, t: int) -> FrameState:
         frame_state = FrameState.kFrameStateInvalid
+        # A final call can flush frames that lack a full 25 ms window (streams under 45 ms) and
+        # so a decibel value; score them as silence, as FsmnVADStreaming.GetFrameState does.
+        if t >= len(self.decibel):
+            return FrameState.kFrameStateSil
         cur_decibel = self.decibel[t]
         cur_snr = cur_decibel - self.noise_average_decibel
         # for each frame, calc log posterior probability of each state

runtime/python/onnxruntime/funasr_onnx/utils/frontend.py

@@ -361,9 +361,11 @@ def extract_fbank(
                     self.lfr_splice_cache[i] = np.concatenate(
                         (self.lfr_splice_cache[i], feats[i]), axis=0
                     )
-                return np.empty(0, dtype=np.float32), feats_lengths
+                # No LFR frame yet; a final call still resets the caches below.
+                feats = np.empty(0, dtype=np.float32)
         else:
-            if is_final:
+            # A stream without a full 25 ms window has no fbank frame to flush.
+            if is_final and self.lfr_splice_cache:
                 self.waveforms = (
                     waveforms if self.reserve_waveforms is None else self.reserve_waveforms
                 )

runtime/python/onnxruntime/funasr_onnx/vad_bin.py

@@ -289,11 +289,11 @@ def __call__(self, audio_in: np.ndarray, **kwargs) -> List:
         is_final = param_dict.get("is_final", False)
         frontend: WavFrontendOnline = param_dict.get("frontend", WavFrontendOnline(cmvn_file=self.cmvn_file, **self.config["frontend_conf"]))
         feats, feats_len = self.extract_feat(frontend=frontend, waveforms=waveforms, is_final=is_final)
+        vad_scorer = param_dict.get("vad_scorer", E2EVadModel(self.config["model_conf"]))
         segments = []
         if feats.size != 0:
             in_cache = param_dict.get("in_cache", list())
             in_cache = self.prepare_cache(in_cache)
-            vad_scorer = param_dict.get("vad_scorer", E2EVadModel(self.config["model_conf"]))
             try:
                 inputs = [feats]
                 inputs.extend(in_cache)
@@ -309,6 +309,9 @@ def __call__(self, audio_in: np.ndarray, **kwargs) -> List:
                 logging.warning("input wav is silence or noise")
                 segments = []
         param_dict.update({"frontend": frontend, "vad_scorer": vad_scorer})
+        if is_final:
+            # Start the next stream with a fresh FSMN cache, as init_cache() and Reset() do.
+            param_dict["in_cache"] = []
         return segments
 
     def load_data(self, wav_content: Union[str, np.ndarray, List[str]], fs: int = None) -> List:

tests/test_funasr_onnx_vad_online_short_chunks.py

@@ -0,0 +1,129 @@
+"""Run against an installed funasr-onnx package."""
+
+import importlib.util
+import itertools
+import unittest
+
+import numpy as np
+
+if importlib.util.find_spec("funasr_onnx") is None:
+    raise unittest.SkipTest("Install the standalone funasr-onnx distribution first.")
+
+from funasr_onnx.vad_bin import Fsmn_vad_online
+
+# frontend_conf of funasr/models/fsmn_vad_streaming/template.yaml
+FRONTEND_CONF = {
+    "fs": 16000,
+    "window": "hamming",
+    "n_mels": 80,
+    "frame_length": 25,
+    "frame_shift": 10,
+    "dither": 0.0,
+    "lfr_m": 5,
+    "lfr_n": 1,
+}
+
+# Segments (ms) of make_audio() streamed in 1600-sample chunks, the step used by the demo.
+EXPECTED = [[[[220, -1]]], [[[-1, 1660]]], [[[2220, -1]]], [[[-1, 3360]]]]
+
+
+def fake_infer(inputs):
+    # Score loud frames as speech instead of running the ONNX session.
+    sil = np.where(inputs[0][0].mean(-1) > 5.0, 0.05, 0.95)
+    scores = np.stack((sil, 1.0 - sil), axis=-1)[None].astype(np.float32)
+    # Change the cache so a cache carried into the next stream shows up.
+    return scores, [cache + 1.0 for cache in inputs[1:]]
+
+
+def make_model():
+    # Keep __call__, WavFrontendOnline and E2EVadModel real; no model files.
+    model = Fsmn_vad_online.__new__(Fsmn_vad_online)
+    model.cmvn_file = None
+    # E2EVadModel's defaults equal the template's model_conf.
+    model.config = {"frontend_conf": FRONTEND_CONF, "model_conf": {}}
+    model.max_end_sil = 800
+    model.encoder_conf = {"fsmn_layers": 1, "proj_dim": 2, "lorder": 2}
+    model.infer = fake_infer
+    return model
+
+
+def make_audio():
+    # 4 s at 16 kHz: noise at 0.5-1.5 s and 2.5-3.2 s, silence elsewhere.
+    rng = np.random.default_rng(0)
+    audio = np.zeros(64000, dtype=np.float32)
+    for start, end in ((8000, 24000), (40000, 51200)):
+        audio[start:end] = rng.uniform(-0.3, 0.3, end - start)
+    return audio
+
+
+def stream(chunk_sizes, param_dict=None):
+    # Feed the audio in chunks, cycling through chunk_sizes; the last chunk is final.
+    model = make_model()
+    audio = make_audio()
+    if param_dict is None:
+        param_dict = {"in_cache": []}
+    results = []
+    offset = 0
+    for size in itertools.cycle(chunk_sizes):
+        chunk = audio[offset : offset + size]
+        offset += size
+        param_dict["is_final"] = offset >= len(audio)
+        segments = model(audio_in=chunk, param_dict=param_dict)
+        if segments:
+            results.append(segments)
+        if param_dict["is_final"]:
+            return results
+
+
+def short_stream(sizes, param_dict):
+    # Send silent chunks of these sizes as a whole stream; the last call is final.
+    model = make_model()
+    results = []
+    for i, size in enumerate(sizes):
+        param_dict["is_final"] = i == len(sizes) - 1
+        results.append(model(audio_in=np.zeros(size, dtype=np.float32), param_dict=param_dict))
+    return results
+
+
+class FsmnVadOnlineShortChunkTest(unittest.TestCase):
+    def test_demo_chunk_size(self):
+        self.assertEqual(stream([1600]), EXPECTED)
+
+    def test_chunks_shorter_than_first_lfr_frame(self):
+        # The first LFR frame needs 720 samples (45 ms), so these first calls yield no features.
+        for size in (160, 320, 480):
+            with self.subTest(size=size):
+                self.assertEqual(stream([size]), EXPECTED)
+
+    def test_packets_that_complete_no_fbank_frame(self):
+        # The first 1600-sample chunk leaves 320 samples for the next 400-sample fbank window,
+        # so a 40-sample packet, or the second of two 80-sample packets, completes no frame.
+        for chunk_sizes in ([1600, 40], [1600, 80, 80]):
+            with self.subTest(chunk_sizes=chunk_sizes):
+                self.assertEqual(stream(chunk_sizes), EXPECTED)
+
+    def test_short_first_and_final_call(self):
+        # Under 400 samples (25 ms) the stream has no fbank frame to flush.
+        for size in (0, 160, 320, 399):
+            with self.subTest(size=size):
+                self.assertEqual(short_stream([size], {"in_cache": []}), [[]])
+
+    def test_empty_final_call_after_partial_chunk(self):
+        # 400-719 samples give one or two fbank frames but no LFR frame. The empty final call
+        # flushes them with less than one 25 ms window of waveform, so they have no decibel value.
+        for size in (400, 480, 560, 719):
+            with self.subTest(size=size):
+                self.assertEqual(short_stream([size, 0], {"in_cache": []}), [[], []])
+
+    def test_param_dict_reuse_after_short_stream(self):
+        # The final call of a short stream resets the frontend, scorer and cache for the next stream.
+        for sizes in ([320], [480, 0], [480]):
+            with self.subTest(sizes=sizes):
+                param_dict = {"in_cache": []}
+                self.assertEqual(short_stream(sizes, param_dict), [[]] * len(sizes))
+                self.assertEqual(param_dict["in_cache"], [])
+                self.assertEqual(stream([1600], param_dict), EXPECTED)
+
+
+if __name__ == "__main__":
+    unittest.main()