fix(asr): end the final VAD speech segment at the last frame's end

nemo/collections/asr/parts/utils/vad_utils.py

@@ -558,7 +558,6 @@ def binarization(sequence: torch.Tensor, per_args: Dict[str, float]) -> torch.Te
 
     speech = False
     start = 0.0
-    i = 0
 
     speech_segments = torch.empty(0)
 
@@ -585,7 +584,9 @@ def binarization(sequence: torch.Tensor, per_args: Dict[str, float]) -> torch.Te
 
     # if it's speech at the end, add final segment
     if speech:
-        new_seg = torch.tensor([max(0, start - pad_onset), i * frame_length_in_sec + pad_offset]).unsqueeze(0)
+        # The last frame is active, so the segment ends where that frame ends, not where it starts.
+        seg_end = len(sequence) * frame_length_in_sec + pad_offset
+        new_seg = torch.tensor([max(0, start - pad_onset), seg_end]).unsqueeze(0)
         speech_segments = torch.cat((speech_segments, new_seg), 0)
 
     # Merge the overlapped speech segments due to padding

tests/collections/speaker_tasks/utils/test_vad_utils_speaker.py

@@ -19,6 +19,7 @@
 
 from nemo.collections.asr.parts.utils.vad_utils import (
     align_labels_to_frames,
+    binarization,
     binarization_vectorized,
     convert_labels_to_speech_segments,
     frame_vad_construct_supervisions_per_file,
@@ -195,6 +196,30 @@ def test_binarization_vectorized_hysteresis(self, predictions, onset, offset, fr
 
         torch.testing.assert_close(segments, torch.tensor(expected))
 
+    @pytest.mark.parametrize(
+        ("predictions", "expected"),
+        [
+            pytest.param([0.6, 0.6], [[0.0, 0.16]], id="speech-runs-to-end-of-audio"),
+            pytest.param([0.1, 0.1, 0.9], [[0.16, 0.24]], id="lone-active-final-frame"),
+            pytest.param([0.9, 0.9, 0.1, 0.9, 0.9], [[0.0, 0.16], [0.24, 0.4]], id="trailing-matches-interior"),
+        ],
+    )
+    @pytest.mark.unit
+    def test_binarization_final_active_frame_uses_full_duration(self, predictions, expected):
+        per_args = {
+            'onset': 0.5,
+            'offset': 0.5,
+            'pad_onset': 0.0,
+            'pad_offset': 0.0,
+            'frame_length_in_sec': 0.08,
+        }
+        sequence = torch.tensor(predictions)
+
+        segments = binarization(sequence, per_args)
+
+        torch.testing.assert_close(segments, torch.tensor(expected))
+        torch.testing.assert_close(segments, binarization_vectorized(sequence, per_args))
+
     @pytest.mark.parametrize(
         "predictions",
         [