|
|
from diart import SpeakerDiarization
|
|
|
from diart.inference import StreamingInference
|
|
|
from diart.sources import AudioSource
|
|
|
from rx.subject import Subject
|
|
|
import threading
|
|
|
import numpy as np
|
|
|
import asyncio
|
|
|
|
|
|
class WebSocketAudioSource(AudioSource):
|
|
|
"""
|
|
|
Simple custom AudioSource that blocks in read()
|
|
|
until close() is called.
|
|
|
push_audio() is used to inject new PCM chunks.
|
|
|
"""
|
|
|
def __init__(self, uri: str = "websocket", sample_rate: int = 16000):
|
|
|
super().__init__(uri, sample_rate)
|
|
|
self._close_event = threading.Event()
|
|
|
self._closed = False
|
|
|
|
|
|
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):
|
|
|
chunk = np.expand_dims(chunk, axis=0)
|
|
|
if not self._closed:
|
|
|
self.stream.on_next(chunk)
|
|
|
|
|
|
|
|
|
def create_pipeline(SAMPLE_RATE):
|
|
|
diar_pipeline = SpeakerDiarization()
|
|
|
ws_source = WebSocketAudioSource(uri="websocket_source", sample_rate=SAMPLE_RATE)
|
|
|
inference = StreamingInference(
|
|
|
pipeline=diar_pipeline,
|
|
|
source=ws_source,
|
|
|
do_plot=False,
|
|
|
show_progress=False,
|
|
|
)
|
|
|
return inference, ws_source
|
|
|
|
|
|
|
|
|
def init_diart(SAMPLE_RATE):
|
|
|
inference, ws_source = create_pipeline(SAMPLE_RATE)
|
|
|
|
|
|
def diar_hook(result):
|
|
|
"""
|
|
|
Hook called each time Diart processes a chunk.
|
|
|
result is (annotation, audio).
|
|
|
We store the label of the last segment in 'current_speaker'.
|
|
|
"""
|
|
|
global l_speakers
|
|
|
l_speakers = []
|
|
|
annotation, audio = result
|
|
|
for speaker in annotation._labels:
|
|
|
segments_beg = annotation._labels[speaker].segments_boundaries_[0]
|
|
|
segments_end = annotation._labels[speaker].segments_boundaries_[-1]
|
|
|
asyncio.create_task(
|
|
|
l_speakers_queue.put({"speaker": speaker, "beg": segments_beg, "end": segments_end})
|
|
|
)
|
|
|
|
|
|
l_speakers_queue = asyncio.Queue()
|
|
|
inference.attach_hooks(diar_hook)
|
|
|
|
|
|
|
|
|
loop = asyncio.get_event_loop()
|
|
|
diar_future = loop.run_in_executor(None, inference)
|
|
|
return inference, l_speakers_queue, ws_source
|
|
|
|
|
|
|
|
|
class DiartDiarization():
|
|
|
def __init__(self, SAMPLE_RATE):
|
|
|
self.inference, self.l_speakers_queue, self.ws_source = init_diart(SAMPLE_RATE)
|
|
|
self.segment_speakers = []
|
|
|
|
|
|
async def diarize(self, pcm_array):
|
|
|
self.ws_source.push_audio(pcm_array)
|
|
|
self.segment_speakers = []
|
|
|
while not self.l_speakers_queue.empty():
|
|
|
self.segment_speakers.append(await self.l_speakers_queue.get())
|
|
|
|
|
|
def close(self):
|
|
|
self.ws_source.close()
|
|
|
|
|
|
|
|
|
def assign_speakers_to_chunks(self, chunks):
|
|
|
"""
|
|
|
Go through each chunk and see which speaker(s) overlap
|
|
|
that chunk's time range in the Diart annotation.
|
|
|
Then store the speaker label(s) (or choose the most overlapping).
|
|
|
This modifies `chunks` in-place or returns a new list with assigned speakers.
|
|
|
"""
|
|
|
if not self.segment_speakers:
|
|
|
return chunks
|
|
|
|
|
|
for segment in self.segment_speakers:
|
|
|
seg_beg = segment["beg"]
|
|
|
seg_end = segment["end"]
|
|
|
speaker = segment["speaker"]
|
|
|
for ch in chunks:
|
|
|
if seg_end <= ch["beg"] or seg_beg >= ch["end"]:
|
|
|
continue
|
|
|
|
|
|
ch["speaker"] = speaker
|
|
|
|
|
|
return chunks
|
|
|
|
