first speaker is "0" no more None

This commit is contained in:
Quentin Fuxa 2025-01-19 19:40:09 +01:00
parent 9bdb92e923
commit 5523b51fd7
3 changed files with 139 additions and 70 deletions

View file

@ -57,9 +57,10 @@ def init_diart(SAMPLE_RATE):
l_speakers = [] l_speakers = []
annotation, audio = result annotation, audio = result
for speaker in annotation._labels: for speaker in annotation._labels:
segment = annotation._labels[speaker].__str__() segments_beg = annotation._labels[speaker].segments_boundaries_[0]
segments_end = annotation._labels[speaker].segments_boundaries_[-1]
asyncio.create_task( asyncio.create_task(
l_speakers_queue.put({"speaker": speaker, "segment": segment}) l_speakers_queue.put({"speaker": speaker, "beg": segments_beg, "end": segments_end})
) )
l_speakers_queue = asyncio.Queue() l_speakers_queue = asyncio.Queue()
@ -74,13 +75,36 @@ def init_diart(SAMPLE_RATE):
class DiartDiarization(): class DiartDiarization():
def __init__(self, SAMPLE_RATE): def __init__(self, SAMPLE_RATE):
self.inference, self.l_speakers_queue, self.ws_source = init_diart(SAMPLE_RATE) self.inference, self.l_speakers_queue, self.ws_source = init_diart(SAMPLE_RATE)
self.segment_speakers = []
async def get_speakers(self, pcm_array): async def diarize(self, pcm_array):
self.ws_source.push_audio(pcm_array) self.ws_source.push_audio(pcm_array)
speakers = [] self.segment_speakers = []
while not self.l_speakers_queue.empty(): while not self.l_speakers_queue.empty():
speakers.append(await self.l_speakers_queue.get()) self.segment_speakers.append(await self.l_speakers_queue.get())
return speakers
def close(self): def close(self):
self.ws_source.close() 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
# We have overlap. Let's just pick the speaker (could be more precise in a more complex implementation)
ch["speaker"] = speaker
return chunks

View file

@ -7,8 +7,8 @@
<style> <style>
body { body {
font-family: 'Inter', sans-serif; font-family: 'Inter', sans-serif;
text-align: center;
margin: 20px; margin: 20px;
text-align: center;
} }
#recordButton { #recordButton {
width: 80px; width: 80px;
@ -28,18 +28,10 @@
#recordButton:active { #recordButton:active {
transform: scale(0.95); transform: scale(0.95);
} }
#transcriptions { #status {
margin-top: 20px; margin-top: 20px;
font-size: 18px; font-size: 16px;
text-align: left; color: #333;
}
.transcription {
display: inline;
color: black;
}
.buffer {
display: inline;
color: rgb(197, 197, 197);
} }
.settings-container { .settings-container {
display: flex; display: flex;
@ -73,9 +65,29 @@
label { label {
font-size: 14px; font-size: 14px;
} }
/* Speaker-labeled transcript area */
#linesTranscript {
margin: 20px auto;
max-width: 600px;
text-align: left;
font-size: 16px;
}
#linesTranscript p {
margin: 5px 0;
}
#linesTranscript strong {
color: #333;
}
/* Grey buffer styling */
.buffer {
color: rgb(180, 180, 180);
font-style: italic;
margin-left: 4px;
}
</style> </style>
</head> </head>
<body> <body>
<div class="settings-container"> <div class="settings-container">
<button id="recordButton">🎙️</button> <button id="recordButton">🎙️</button>
<div class="settings"> <div class="settings">
@ -96,9 +108,11 @@
</div> </div>
</div> </div>
</div> </div>
<p id="status"></p> <p id="status"></p>
<div id="transcriptions"></div> <!-- Speaker-labeled transcript -->
<div id="linesTranscript"></div>
<script> <script>
let isRecording = false; let isRecording = false;
@ -106,89 +120,97 @@
let recorder = null; let recorder = null;
let chunkDuration = 1000; let chunkDuration = 1000;
let websocketUrl = "ws://localhost:8000/asr"; let websocketUrl = "ws://localhost:8000/asr";
// Tracks whether the user voluntarily closed the WebSocket
let userClosing = false; let userClosing = false;
const statusText = document.getElementById("status"); const statusText = document.getElementById("status");
const recordButton = document.getElementById("recordButton"); const recordButton = document.getElementById("recordButton");
const chunkSelector = document.getElementById("chunkSelector"); const chunkSelector = document.getElementById("chunkSelector");
const websocketInput = document.getElementById("websocketInput"); const websocketInput = document.getElementById("websocketInput");
const transcriptionsDiv = document.getElementById("transcriptions"); const linesTranscriptDiv = document.getElementById("linesTranscript");
let fullTranscription = ""; // Store confirmed transcription
// Update chunk duration based on the selector
chunkSelector.addEventListener("change", () => { chunkSelector.addEventListener("change", () => {
chunkDuration = parseInt(chunkSelector.value); chunkDuration = parseInt(chunkSelector.value);
}); });
// Update WebSocket URL dynamically, with some basic checks
websocketInput.addEventListener("change", () => { websocketInput.addEventListener("change", () => {
const urlValue = websocketInput.value.trim(); const urlValue = websocketInput.value.trim();
// Quick check to see if it starts with ws:// or wss://
if (!urlValue.startsWith("ws://") && !urlValue.startsWith("wss://")) { if (!urlValue.startsWith("ws://") && !urlValue.startsWith("wss://")) {
statusText.textContent = statusText.textContent = "Invalid WebSocket URL (must start with ws:// or wss://)";
"Invalid WebSocket URL. It should start with ws:// or wss://";
return; return;
} }
websocketUrl = urlValue; websocketUrl = urlValue;
statusText.textContent = "WebSocket URL updated. Ready to connect."; statusText.textContent = "WebSocket URL updated. Ready to connect.";
}); });
/**
* Opens webSocket connection.
* returns a Promise that resolves when the connection is open.
* rejects if there was an error.
*/
function setupWebSocket() { function setupWebSocket() {
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
try { try {
websocket = new WebSocket(websocketUrl); websocket = new WebSocket(websocketUrl);
} catch (error) { } catch (error) {
statusText.textContent = statusText.textContent = "Invalid WebSocket URL. Please check and try again.";
"Invalid WebSocket URL. Please check the URL and try again.";
reject(error); reject(error);
return; return;
} }
websocket.onopen = () => { websocket.onopen = () => {
statusText.textContent = "Connected to server"; statusText.textContent = "Connected to server.";
resolve(); resolve();
}; };
websocket.onclose = (event) => { websocket.onclose = () => {
// If we manually closed it, we say so
if (userClosing) { if (userClosing) {
statusText.textContent = "WebSocket closed by user."; statusText.textContent = "WebSocket closed by user.";
} else { } else {
statusText.textContent = "Disconnected from the websocket server. If this is the first launch, the model may be downloading in the backend. Check the API logs for more information."; statusText.textContent =
"Disconnected from the WebSocket server. (Check logs if model is loading.)";
} }
userClosing = false; userClosing = false;
}; };
websocket.onerror = () => { websocket.onerror = () => {
statusText.textContent = "Error connecting to WebSocket"; statusText.textContent = "Error connecting to WebSocket.";
reject(new Error("Error connecting to WebSocket")); reject(new Error("Error connecting to WebSocket"));
}; };
// Handle messages from server
websocket.onmessage = (event) => { websocket.onmessage = (event) => {
const data = JSON.parse(event.data); const data = JSON.parse(event.data);
const { transcription, buffer } = data; /*
The server might send:
// Update confirmed transcription {
fullTranscription += transcription; "lines": [
{"speaker": 0, "text": "Hello."},
// Update the transcription display {"speaker": 1, "text": "Bonjour."},
transcriptionsDiv.innerHTML = ` ...
<span class="transcription">${fullTranscription}</span> ],
<span class="buffer">${buffer}</span> "buffer": "..."
`; }
*/
const { lines = [], buffer = "" } = data;
renderLinesWithBuffer(lines, buffer);
}; };
}); });
} }
function renderLinesWithBuffer(lines, buffer) {
// Clears if no lines
if (!Array.isArray(lines) || lines.length === 0) {
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 textContent = item.text;
if (idx === lines.length - 1 && buffer) {
textContent += `<span class="buffer">${buffer}</span>`;
}
return `<p><strong>Speaker ${item.speaker}:</strong> ${textContent}</p>`;
}).join("");
linesTranscriptDiv.innerHTML = linesHtml;
}
async function startRecording() { async function startRecording() {
try { try {
const stream = await navigator.mediaDevices.getUserMedia({ audio: true }); const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
@ -202,22 +224,18 @@
isRecording = true; isRecording = true;
updateUI(); updateUI();
} catch (err) { } catch (err) {
statusText.textContent = statusText.textContent = "Error accessing microphone. Please allow microphone access.";
"Error accessing microphone. Please allow microphone access.";
} }
} }
function stopRecording() { function stopRecording() {
userClosing = true; userClosing = true;
// Stop the recorder if it exists
if (recorder) { if (recorder) {
recorder.stop(); recorder.stop();
recorder = null; recorder = null;
} }
isRecording = false; isRecording = false;
// Close the websocket if it exists
if (websocket) { if (websocket) {
websocket.close(); websocket.close();
websocket = null; websocket = null;
@ -228,15 +246,12 @@
async function toggleRecording() { async function toggleRecording() {
if (!isRecording) { if (!isRecording) {
fullTranscription = ""; linesTranscriptDiv.innerHTML = "";
transcriptionsDiv.innerHTML = "";
try { try {
await setupWebSocket(); await setupWebSocket();
await startRecording(); await startRecording();
} catch (err) { } catch (err) {
statusText.textContent = statusText.textContent = "Could not connect to WebSocket or access mic. Aborted.";
"Could not connect to WebSocket or access mic. Recording aborted.";
} }
} else { } else {
stopRecording(); stopRecording();
@ -245,9 +260,7 @@
function updateUI() { function updateUI() {
recordButton.classList.toggle("recording", isRecording); recordButton.classList.toggle("recording", isRecording);
statusText.textContent = isRecording statusText.textContent = isRecording ? "Recording..." : "Click to start transcription";
? "Recording..."
: "Click to start transcription";
} }
recordButton.addEventListener("click", toggleRecording); recordButton.addEventListener("click", toggleRecording);

View file

@ -90,6 +90,7 @@ async def start_ffmpeg_decoder():
return process return process
@app.websocket("/asr") @app.websocket("/asr")
async def websocket_endpoint(websocket: WebSocket): async def websocket_endpoint(websocket: WebSocket):
await websocket.accept() await websocket.accept()
@ -110,6 +111,9 @@ async def websocket_endpoint(websocket: WebSocket):
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
full_transcription = "" full_transcription = ""
beg = time() beg = time()
chunk_history = [] # Will store dicts: {beg, end, text, speaker}
while True: while True:
try: try:
elapsed_time = int(time() - beg) elapsed_time = int(time() - beg)
@ -137,8 +141,17 @@ async def websocket_endpoint(websocket: WebSocket):
) )
pcm_buffer = bytearray() pcm_buffer = bytearray()
online.insert_audio_chunk(pcm_array) online.insert_audio_chunk(pcm_array)
transcription = online.process_iter()[2] beg_trans, end_trans, trans = online.process_iter()
full_transcription += transcription
if trans:
chunk_history.append({
"beg": beg_trans,
"end": end_trans,
"text": trans,
"speaker": "0"
})
full_transcription += trans
if args.vac: if args.vac:
buffer = online.online.to_flush( buffer = online.online.to_flush(
online.online.transcript_buffer.buffer online.online.transcript_buffer.buffer
@ -151,11 +164,30 @@ async def websocket_endpoint(websocket: WebSocket):
buffer in full_transcription buffer in full_transcription
): # With VAC, the buffer is not updated until the next chunk is processed ): # With VAC, the buffer is not updated until the next chunk is processed
buffer = "" buffer = ""
response = {"transcription": transcription, "buffer": buffer}
lines = [
{
"speaker": "0",
"text": "",
}
]
if args.diarization: if args.diarization:
speakers = await diarization.get_speakers(pcm_array) await diarization.diarize(pcm_array)
response["speakers"] = speakers diarization.assign_speakers_to_chunks(chunk_history)
for ch in chunk_history:
if args.diarization and ch["speaker"] and ch["speaker"][-1] != lines[-1]["speaker"]:
lines.append(
{
"speaker": ch["speaker"][-1],
"text": ch['text'],
}
)
else:
lines[-1]["text"] += ch['text']
response = {"lines": lines, "buffer": buffer}
await websocket.send_json(response) await websocket.send_json(response)
except Exception as e: except Exception as e: