admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #51 from QuentinFuxa/diart_integration_improvements

Browse source 
Diart integration improvements
This commit is contained in:
Quentin Fuxa 2025-02-19 11:38:59 +01:00 committed by GitHub
commit 450c93fef8
3 changed files with 121 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

76
src/diarization/diarization_online.py
View file

@ -5,6 +5,11 @@ from rx.subject import Subject
import threading
import numpy as np
import asyncio
import re
def extract_number(s):
match = re.search(r'\d+', s)
return int(match.group()) if match else None
class WebSocketAudioSource(AudioSource):
"""
@ -44,37 +49,48 @@ def create_pipeline(SAMPLE_RATE):
return inference, ws_source
def init_diart(SAMPLE_RATE):
inference, ws_source = create_pipeline(SAMPLE_RATE)
def init_diart(SAMPLE_RATE, diar_instance):
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,
)
l_speakers_queue = asyncio.Queue()
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'.
For each detected speaker segment, push its info to the queue and update processed_time.
"""
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})
)
if annotation._labels:
for speaker in annotation._labels:
segments_beg = annotation._labels[speaker].segments_boundaries_[0]
segments_end = annotation._labels[speaker].segments_boundaries_[-1]
if segments_end > diar_instance.processed_time:
diar_instance.processed_time = segments_end
asyncio.create_task(
l_speakers_queue.put({"speaker": speaker, "beg": segments_beg, "end": segments_end})
)
else:
audio_duration = audio.extent.end
if audio_duration > diar_instance.processed_time:
diar_instance.processed_time = audio_duration
l_speakers_queue = asyncio.Queue()
inference.attach_hooks(diar_hook)
# Launch Diart in a background thread
loop = asyncio.get_event_loop()
diar_future = loop.run_in_executor(None, inference)
return inference, l_speakers_queue, ws_source
class DiartDiarization():
class DiartDiarization:
def __init__(self, SAMPLE_RATE):
self.inference, self.l_speakers_queue, self.ws_source = init_diart(SAMPLE_RATE)
self.processed_time = 0
self.inference, self.l_speakers_queue, self.ws_source = init_diart(SAMPLE_RATE, self)
self.segment_speakers = []
async def diarize(self, pcm_array):
@ -82,20 +98,21 @@ class DiartDiarization():
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.
For each chunk (a dict with keys "beg" and "end"), assign a speaker label.
- If a chunk overlaps with a detected speaker segment, assign that label.
- If the chunk's end time is within the processed time and no speaker was assigned,
mark it as "No speaker".
- If the chunk's time hasn't been fully processed yet, leave it (or mark as "Processing").
"""
if not self.segment_speakers:
return chunks
for ch in chunks:
ch["speaker"] = ch.get("speaker", -1)
for segment in self.segment_speakers:
seg_beg = segment["beg"]
@ -104,7 +121,10 @@ class DiartDiarization():
for ch in chunks:
if seg_end <= ch["beg"] or seg_beg >= ch["end"]:
continue
# We have overlap. Let's just pick the speaker (could be more precise in a more complex implementation)
ch["speaker"] = speaker
ch["speaker"] = extract_number(speaker) + 1
if self.processed_time > 0:
for ch in chunks:
if ch["end"] <= self.processed_time and ch["speaker"] == -1:
ch["speaker"] = -2
return chunks
return chunks

24
src/web/live_transcription.html
View file

@ -179,8 +179,9 @@
The server might send:
{
"lines": [
{"speaker": 0, "text": "Hello."},
{"speaker": 1, "text": "Bonjour."},
{"speaker": 0, "text": "Hello.", "beg": "00:00", "end": "00:01"},
{"speaker": -2, "text": "Hi, no speaker here.", "beg": "00:01", "end": "00:02"},
{"speaker": -1, "text": "...", "beg": "00:02", "end": "00:03" },
...
],
"buffer": "..."
@ -198,14 +199,27 @@
linesTranscriptDiv.innerHTML = "";
return;
}
// Build the HTML
// The buffer is appended to the last line if it's non-empty
const linesHtml = lines.map((item, idx) => {
let speakerLabel = "";
if (item.speaker === -2) {
speakerLabel = "No speaker";
} else if (item.speaker !== -1) {
speakerLabel = `Speaker ${item.speaker}`;
}
let timeInfo = "";
if (item.beg !== undefined && item.end !== undefined) {
timeInfo = ` [${item.beg}, ${item.end}]`;
}
let textContent = item.text;
if (idx === lines.length - 1 && buffer) {
textContent += `<span class="buffer">${buffer}</span>`;
}
return `<p><strong>Speaker ${item.speaker}:</strong> ${textContent}</p>`;
return speakerLabel
? `<p><strong>${speakerLabel}${timeInfo}</strong> ${textContent}</p>`
: `<p>${textContent}</p>`;
}).join("");
linesTranscriptDiv.innerHTML = linesHtml;

84
whisper_fastapi_online_server.py
View file

@ -3,7 +3,7 @@ import argparse
import asyncio
import numpy as np
import ffmpeg
from time import time
from time import time, sleep
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
@ -12,9 +12,12 @@ from fastapi.middleware.cors import CORSMiddleware
from src.whisper_streaming.whisper_online import backend_factory, online_factory, add_shared_args
import subprocess
import math
import logging
from datetime import timedelta
def format_time(seconds):
return str(timedelta(seconds=int(seconds)))
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@ -48,6 +51,12 @@ parser.add_argument(
help="Whether to enable speaker diarization.",
)
parser.add_argument(
"--transcription",
type=bool,
default=True,
help="To disable to only see live diarization results.",
)
add_shared_args(parser)
args = parser.parse_args()
@ -68,7 +77,10 @@ if args.diarization:
@asynccontextmanager
async def lifespan(app: FastAPI):
global asr, tokenizer
asr, tokenizer = backend_factory(args)
if args.transcription:
asr, tokenizer = backend_factory(args)
else:
asr, tokenizer = None, None
yield
app = FastAPI(lifespan=lifespan)
@ -117,7 +129,7 @@ async def websocket_endpoint(websocket: WebSocket):
ffmpeg_process = None
pcm_buffer = bytearray()
online = online_factory(args, asr, tokenizer)
online = online_factory(args, asr, tokenizer) if args.transcription else None
diarization = DiartDiarization(SAMPLE_RATE) if args.diarization else None
async def restart_ffmpeg():
@ -130,7 +142,7 @@ async def websocket_endpoint(websocket: WebSocket):
logger.warning(f"Error killing FFmpeg process: {e}")
ffmpeg_process = await start_ffmpeg_decoder()
pcm_buffer = bytearray()
online = online_factory(args, asr, tokenizer)
online = online_factory(args, asr, tokenizer) if args.transcription else None
if args.diarization:
diarization = DiartDiarization(SAMPLE_RATE)
logger.info("FFmpeg process started.")
@ -142,7 +154,7 @@ async def websocket_endpoint(websocket: WebSocket):
loop = asyncio.get_event_loop()
full_transcription = ""
beg = time()
beg_loop = time()
chunk_history = [] # Will store dicts: {beg, end, text, speaker}
while True:
@ -184,45 +196,57 @@ async def websocket_endpoint(websocket: WebSocket):
/ 32768.0
)
pcm_buffer = pcm_buffer[MAX_BYTES_PER_SEC:]
logger.info(f"{len(online.audio_buffer) / online.SAMPLING_RATE} seconds of audio will be processed by the model.")
online.insert_audio_chunk(pcm_array)
transcription = online.process_iter()
if transcription:
if args.transcription:
logger.info(f"{len(online.audio_buffer) / online.SAMPLING_RATE} seconds of audio will be processed by the model.")
online.insert_audio_chunk(pcm_array)
transcription = online.process_iter()
if transcription.start:
chunk_history.append({
"beg": transcription.start,
"end": transcription.end,
"text": transcription.text,
})
full_transcription += transcription.text if transcription else ""
buffer = online.get_buffer()
if buffer in full_transcription: # With VAC, the buffer is not updated until the next chunk is processed
buffer = ""
else:
chunk_history.append({
"beg": transcription.start,
"end": transcription.end,
"text": transcription.text,
"speaker": "0"
"beg": time() - beg_loop,
"end": time() - beg_loop + 0.1,
"text": '',
})
sleep(0.1)
buffer = ''
full_transcription += transcription.text if transcription else ""
buffer = online.get_buffer()
if buffer in full_transcription: # With VAC, the buffer is not updated until the next chunk is processed
buffer = ""
lines = [
{
"speaker": "0",
"text": "",
}
]
if args.diarization:
await diarization.diarize(pcm_array)
diarization.assign_speakers_to_chunks(chunk_history)
current_speaker = -1
lines = [{
"beg": 0,
"end": 0,
"speaker": current_speaker,
"text": ""
}]
for ch in chunk_history:
if args.diarization and ch["speaker"] and ch["speaker"][-1] != lines[-1]["speaker"]:
if args.diarization and ch["speaker"] and ch["speaker"] != current_speaker:
new_speaker = ch["speaker"]
lines.append(
{
"speaker": ch["speaker"][-1],
"text": ch['text']
"speaker": new_speaker,
"text": ch['text'],
"beg": format_time(ch['beg']),
"end": format_time(ch['end']),
}
)
current_speaker = new_speaker
else:
lines[-1]["text"] += ch['text']
lines[-1]["end"] = format_time(ch['end'])
response = {"lines": lines, "buffer": buffer}
await websocket.send_json(response)