diff --git a/nemo/collections/asr/parts/utils/vad_utils.py b/nemo/collections/asr/parts/utils/vad_utils.py index 3037f2ddec9d..9f3e88680d87 100644 --- a/nemo/collections/asr/parts/utils/vad_utils.py +++ b/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 diff --git a/tests/collections/speaker_tasks/utils/test_vad_utils_speaker.py b/tests/collections/speaker_tasks/utils/test_vad_utils_speaker.py index 29d0dd79585f..f6ba88ca8de0 100644 --- a/tests/collections/speaker_tasks/utils/test_vad_utils_speaker.py +++ b/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", [