update import paths

diarization/diarization_online.py

@@ -0,0 +1,90 @@
+import asyncio
+import re
+import threading
+import numpy as np
+
+from diart import SpeakerDiarization
+from diart.inference import StreamingInference
+from diart.sources import AudioSource
+from timed_objects import SpeakerSegment
+
+def extract_number(s: str) -> int:
+    m = re.search(r'\d+', s)
+    return int(m.group()) if m else None
+
+
+class WebSocketAudioSource(AudioSource):
+    """
+    Custom AudioSource that blocks in read() until close() is called.
+    Use push_audio() to inject PCM chunks.
+    """
+    def __init__(self, uri: str = "websocket", sample_rate: int = 16000):
+        super().__init__(uri, sample_rate)
+        self._closed = False
+        self._close_event = threading.Event()
+
+    def read(self):
+        self._close_event.wait()
+
+    def close(self):
+        if not self._closed:
+            self._closed = True
+            self.stream.on_completed()
+            self._close_event.set()
+
+    def push_audio(self, chunk: np.ndarray):
+        if not self._closed:
+            self.stream.on_next(np.expand_dims(chunk, axis=0))
+
+
+class DiartDiarization:
+    def __init__(self, sample_rate: int):
+        self.processed_time = 0
+        self.segment_speakers = []
+        self.speakers_queue = asyncio.Queue()
+        self.pipeline = SpeakerDiarization()
+        self.source = WebSocketAudioSource(uri="websocket_source", sample_rate=sample_rate)
+        self.inference = StreamingInference(
+            pipeline=self.pipeline,
+            source=self.source,
+            do_plot=False,
+            show_progress=False,
+        )
+        # Attache la fonction hook et démarre l'inférence en arrière-plan.
+        self.inference.attach_hooks(self._diar_hook)
+        asyncio.get_event_loop().run_in_executor(None, self.inference)
+
+    def _diar_hook(self, result):
+        annotation, audio = result
+        if annotation._labels:
+            for speaker, label in annotation._labels.items():
+                start = label.segments_boundaries_[0]
+                end = label.segments_boundaries_[-1]
+                if end > self.processed_time:
+                    self.processed_time = end
+                asyncio.create_task(self.speakers_queue.put(SpeakerSegment(
+                    speaker=speaker,
+                    start=start,
+                    end=end,
+                )))
+        else:
+            dur = audio.extent.end
+            if dur > self.processed_time:
+                self.processed_time = dur
+
+    async def diarize(self, pcm_array: np.ndarray):
+        self.source.push_audio(pcm_array)
+        self.segment_speakers.clear()
+        while not self.speakers_queue.empty():
+            self.segment_speakers.append(await self.speakers_queue.get())
+
+    def close(self):
+        self.source.close()
+
+    def assign_speakers_to_tokens(self, end_attributed_speaker, tokens: list) -> list:
+        for token in tokens:
+            for segment in self.segment_speakers:
+                if not (segment.end <= token.start or segment.start >= token.end):
+                    token.speaker = extract_number(segment.speaker) + 1
+                    end_attributed_speaker = max(token.end, end_attributed_speaker)
+        return end_attributed_speaker