From 6121083549c129aba048f61e6ac1c11a558c3102 Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Wed, 19 Feb 2025 11:21:40 +0100 Subject: [PATCH 1/3] add parameter to disable transcription (only diarization), add time in output --- whisper_fastapi_online_server.py | 84 ++++++++++++++++++++------------ 1 file changed, 54 insertions(+), 30 deletions(-) diff --git a/whisper_fastapi_online_server.py b/whisper_fastapi_online_server.py index 89b010c..e4a1571 100644 --- a/whisper_fastapi_online_server.py +++ b/whisper_fastapi_online_server.py @@ -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) From dc24366580dec71c0452e0f2d6205430c8aac9fb Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Wed, 19 Feb 2025 11:25:59 +0100 Subject: [PATCH 2/3] Number of speakers not anymore limited to 10; a speaker has been created for "being processed" (-1), and another one for no" speaker detected" (-2) --- src/diarization/diarization_online.py | 76 +++++++++++++++++---------- 1 file changed, 48 insertions(+), 28 deletions(-) diff --git a/src/diarization/diarization_online.py b/src/diarization/diarization_online.py index 432e104..8240408 100644 --- a/src/diarization/diarization_online.py +++ b/src/diarization/diarization_online.py @@ -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 \ No newline at end of file From 1ffa2fa224f016f19a3693240cbf1c12d93d463f Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Wed, 19 Feb 2025 11:26:30 +0100 Subject: [PATCH 3/3] add time duration ; display when no speaker is detected --- src/web/live_transcription.html | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/web/live_transcription.html b/src/web/live_transcription.html index 0c6d0ba..af48481 100644 --- a/src/web/live_transcription.html +++ b/src/web/live_transcription.html @@ -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 += `${buffer}`; } - return `

Speaker ${item.speaker}: ${textContent}

`; + + return speakerLabel + ? `

${speakerLabel}${timeInfo} ${textContent}

` + : `

${textContent}

`; }).join(""); linesTranscriptDiv.innerHTML = linesHtml;