Merge branch 'pa-speaker-detector-robustness' into 'main'

ENH: Improve robustness of PASpeakerDetector See merge request heka/medkit!245 changelog: ENH: Improve robustness of PASpeakerDetector
medkit-lib · Dec 13, 2023 · f4b0687 · f4b0687
2 parents 25ed286 + 6f05446
commit f4b0687
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 12 deletions.
diff --git a/medkit/audio/segmentation/pa_speaker_detector.py b/medkit/audio/segmentation/pa_speaker_detector.py
@@ -25,6 +25,11 @@
 from medkit.core.audio import Segment, SegmentationOperation, Span
 
 
+# margin (in seconds) by which a turn segment
+# may overrun the input segment due to imprecision
+_DURATION_MARGIN = 0.1
+
+
 class PASpeakerDetector(SegmentationOperation):
     """Speaker diarization operation relying on `pyannote.audio`
 
@@ -150,12 +155,16 @@ def _detect_turns_in_segment(self, segment: Segment) -> Iterator[Segment]:
         for turn, _, speaker in diarization.itertracks(yield_label=True):
             if turn.duration < self.min_duration:
                 continue
+
             # trim original audio to turn start/end points
-            turn_audio = audio.trim_duration(turn.start, turn.end)
+            # (allow pyannote's turn to be slighty over the total input duration)
+            assert turn.end < audio.duration + _DURATION_MARGIN
+            turn_end = min(turn.end, audio.duration)
+            turn_audio = audio.trim_duration(turn.start, turn_end)
 
             turn_span = Span(
                 start=segment.span.start + turn.start,
-                end=segment.span.start + turn.end,
+                end=segment.span.start + turn_end,
             )
             speaker_attr = Attribute(label="speaker", value=speaker)
             turn_segment = Segment(

diff --git a/tests/unit/audio/metrics/test_transcription_evaluator.py b/tests/unit/audio/metrics/test_transcription_evaluator.py
@@ -24,15 +24,15 @@ def _get_doc():
     turn_seg_1 = Segment(
         label="speech",
         audio=_FULL_AUDIO.trim_duration(start_time=0.0, end_time=2.0),
-        span=Span(start=0.0, end=4.0),
+        span=Span(start=0.0, end=2.0),
         attrs=[Attribute(label="transcription", value="Bonjour ça va bien ?")],
     )
     doc.anns.add(turn_seg_1)
 
     turn_seg_2 = Segment(
         label="speech",
         audio=_FULL_AUDIO.trim_duration(start_time=2.0, end_time=4.0),
-        span=Span(5.0, 6.0),
+        span=Span(2.0, 4.0),
         attrs=[Attribute(label="transcription", value="Ça va et vous ?")],
     )
     doc.anns.add(turn_seg_2)
@@ -45,7 +45,7 @@ def _get_doc():
     "identical": (
         [
             {"start": 0.0, "end": 2.0, "transcription": "Bonjour ça va bien ?"},
-            {"start": 2.0, "end": 4.0, "transcription": "Ça va et vous ?"},
+            {"start": 3.0, "end": 4.0, "transcription": "Ça va et vous ?"},
         ],
         {},
         TranscriptionEvaluatorResult(
@@ -65,7 +65,7 @@ def _get_doc():
     "errors": (
         [
             {"start": 0.0, "end": 2.0, "transcription": "Bonjour ça va ?"},
-            {"start": 2.0, "end": 4.0, "transcription": "Bien et vous ?"},
+            {"start": 3.0, "end": 4.0, "transcription": "Bien et vous ?"},
         ],
         {},
         TranscriptionEvaluatorResult(

diff --git a/tests/unit/audio/preprocessing/test_downmixer.py b/tests/unit/audio/preprocessing/test_downmixer.py
@@ -11,9 +11,11 @@
 
 
 def _get_segment(signal):
-    duration = signal.shape[1] * _SAMPLE_RATE
+    duration = signal.shape[1] / _SAMPLE_RATE
     audio = MemoryAudioBuffer(signal=signal, sample_rate=_SAMPLE_RATE)
-    return Segment(label="raw", span=Span(_SPAN_OFFSET, duration), audio=audio)
+    return Segment(
+        label="raw", span=Span(_SPAN_OFFSET, _SPAN_OFFSET + duration), audio=audio
+    )
 
 
 def _check_downmixed_segment(downmixed_seg, original_seg):

diff --git a/tests/unit/audio/preprocessing/test_power_normalizer.py b/tests/unit/audio/preprocessing/test_power_normalizer.py
@@ -14,9 +14,11 @@
 
 
 def _get_segment(signal):
-    duration = signal.shape[1] * _SAMPLE_RATE
+    duration = signal.shape[1] / _SAMPLE_RATE
     audio = MemoryAudioBuffer(signal=signal, sample_rate=_SAMPLE_RATE)
-    return Segment(label="raw", span=Span(_SPAN_OFFSET, duration), audio=audio)
+    return Segment(
+        label="raw", span=Span(_SPAN_OFFSET, _SPAN_OFFSET + duration), audio=audio
+    )
 
 
 def _check_normalized_segment(normalized_seg, original_seg):

diff --git a/tests/unit/audio/preprocessing/test_resampler.py b/tests/unit/audio/preprocessing/test_resampler.py
@@ -15,9 +15,11 @@
 
 
 def _get_segment(signal, sample_rate):
-    duration = signal.shape[1] * sample_rate
+    duration = signal.shape[1] / sample_rate
     audio = MemoryAudioBuffer(signal=signal, sample_rate=sample_rate)
-    return Segment(label="raw", span=Span(_SPAN_OFFSET, duration), audio=audio)
+    return Segment(
+        label="raw", span=Span(_SPAN_OFFSET, _SPAN_OFFSET + duration), audio=audio
+    )
 
 
 def _check_resampled_segment(resampled_seg, original_seg):