""" Speaker Identification — Gradio app for Hugging Face Spaces Upload a WAV file + a JSON file of transcript segments ([{start, end, text}, ...]) and get back the segments labeled with speaker IDs (Speaker 1, Speaker 2, ...). This mirrors the pipeline from the "ZzzFix_Export_Real_Audio" Colab notebook: 1. Load speechbrain/spkrec-resnet-voxceleb embedding model 2. Slice audio into ~1.5s windows per transcript segment 3. Extract a normalized embedding per window 4. Cluster embeddings (AgglomerativeClustering, cosine, average linkage), picking the k with the best silhouette score 5. Merge consecutive same-speaker windows and align back to text """ import os import json import struct import warnings import tempfile import torch import torch.nn.functional as F import numpy as np import gradio as gr from pydub import AudioSegment from sklearn.cluster import AgglomerativeClustering from sklearn.metrics import silhouette_score warnings.filterwarnings("ignore") torch.set_num_threads(os.cpu_count() or 4) # --------------------------------------------------------------------------- # Load model once at startup # --------------------------------------------------------------------------- from speechbrain.inference import EncoderClassifier # speechbrain >= 1.0 # If you're on an older speechbrain version, use this instead: # from speechbrain.pretrained import EncoderClassifier print("Loading speaker embedding model...") classifier = EncoderClassifier.from_hparams( source="speechbrain/spkrec-resnet-voxceleb", run_opts={"device": "cpu"}, ) classifier.eval() print("Model loaded.") WINDOW = 1.5 # seconds per embedding window MIN_WINDOW = 0.5 # discard trailing windows shorter than this SILENCE_DBFS = -45 # skip near-silent windows PCM_SAMPLE_RATE = 16000 def load_audio_samples(audio_path: str): """Return (samples_float32, sample_rate). pcm.bin (header: int32 n_samples, float32 duration, then float32[n_samples] mono 16kHz samples in -1..1) is read directly — no ffmpeg/pydub decode, which is the main time saver. Falls back to pydub for .wav/other files.""" ext = os.path.splitext(audio_path)[1].lower() if ext in (".bin", ".pcm", ".raw"): with open(audio_path, "rb") as f: header = f.read(8) n_samples, _duration = struct.unpack(" np.ndarray: """Run all window chunks through the model in batches instead of one at a time. This is the single biggest speedup on CPU — one forward pass per batch instead of one forward pass per ~1.5s window.""" all_embs = [] for i in range(0, len(chunks), batch_size): batch = chunks[i:i + batch_size] lengths = [c.shape[1] for c in batch] max_len = max(lengths) padded = torch.zeros(len(batch), max_len) for j, c in enumerate(batch): padded[j, :c.shape[1]] = c[0] wav_lens = torch.tensor([l / max_len for l in lengths]) with torch.no_grad(): embs = classifier.encode_batch(padded, wav_lens) # [B, 1, D] embs = embs.squeeze(1) embs = F.normalize(embs, dim=1) all_embs.append(embs.numpy()) return np.concatenate(all_embs, axis=0) def diarize(audio_path: str, json_path: str): """Run the full pipeline and return (dataframe_rows, status_message).""" try: with open(json_path) as f: segments = json.load(f) except Exception as e: return [], f"Could not read JSON segments file: {e}" if not isinstance(segments, list) or len(segments) == 0: return [], "JSON file must be a non-empty list of {start, end, text} objects." # Load audio — pcm.bin skips ffmpeg/pydub decode entirely (fast path); # .wav/other files still fall back to pydub. try: samples_all, sr = load_audio_samples(audio_path) except Exception as e: return [], f"Could not read audio file: {e}" subsegments, chunks = [], [] for seg in segments: try: t, seg_end = float(seg["start"]), float(seg["end"]) except (KeyError, TypeError, ValueError): continue text = seg.get("text", "") while t < seg_end: s, e = t, min(t + WINDOW, seg_end) if (e - s) < MIN_WINDOW: break start_idx = int(s * sr) end_idx = int(e * sr) chunk_samples = samples_all[start_idx:end_idx] if chunk_samples.size == 0: t += WINDOW continue rms = np.sqrt(np.mean(chunk_samples ** 2)) + 1e-9 dbfs = 20 * np.log10(rms) if dbfs < SILENCE_DBFS: t += WINDOW continue signal = torch.from_numpy(chunk_samples.copy()).unsqueeze(0) subsegments.append({"start": s, "end": e, "text": text}) chunks.append(signal) t += WINDOW if len(chunks) < 2: return [], "Not enough speech windows extracted to cluster speakers (need at least 2)." embeddings = get_embeddings_batched(chunks) # Pick best k by silhouette score, exactly like the notebook best_score, best_labels, best_k = -999, None, 2 max_k = min(6, len(embeddings) - 1) for k in range(2, max_k + 1): try: lbl = AgglomerativeClustering( n_clusters=k, metric="cosine", linkage="average" ).fit_predict(embeddings) if len(set(lbl)) < 2: continue score = silhouette_score(embeddings, lbl, metric="cosine") if score > best_score: best_score, best_labels, best_k = score, lbl, k except Exception: pass if best_labels is None: return [], "Clustering failed — try a longer audio clip with more speech." speaker_labels = [f"Speaker {x + 1}" for x in best_labels] # Merge consecutive windows from the same speaker into readable rows rows = [] cur_speaker = speaker_labels[0] cur_start = subsegments[0]["start"] cur_end = subsegments[0]["end"] cur_text = [subsegments[0]["text"]] for i in range(1, len(subsegments)): if speaker_labels[i] == cur_speaker: cur_end = subsegments[i]["end"] cur_text.append(subsegments[i]["text"]) else: rows.append([ f"{cur_start:.2f}", f"{cur_end:.2f}", cur_speaker, " ".join(t for t in cur_text if t).strip(), ]) cur_speaker = speaker_labels[i] cur_start = subsegments[i]["start"] cur_end = subsegments[i]["end"] cur_text = [subsegments[i]["text"]] rows.append([ f"{cur_start:.2f}", f"{cur_end:.2f}", cur_speaker, " ".join(t for t in cur_text if t).strip(), ]) status = f"Detected {best_k} speaker(s) — silhouette score: {best_score:.4f}" return rows, status def run_pipeline(audio_file, json_file): if audio_file is None or json_file is None: return None, "Please upload both a WAV file and a JSON segments file." rows, status = diarize(audio_file, json_file) return rows, status # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- with gr.Blocks(title="Speaker Identification") as demo: gr.Markdown( """ # 🎙️ Speaker Identification Upload an audio file (.wav) and its matching transcript JSON (a list of `{"start": ..., "end": ..., "text": ...}` objects), and this will cluster the speech into speakers using `speechbrain/spkrec-resnet-voxceleb` embeddings. """ ) with gr.Row(): audio_input = gr.File(label="Audio (pcm.bin or .wav)", file_types=[".bin", ".pcm", ".raw", ".wav"]) json_input = gr.File(label="Transcript JSON", file_types=[".json"]) run_btn = gr.Button("Identify Speakers", variant="primary") status_box = gr.Markdown() output_table = gr.Dataframe( headers=["Start (s)", "End (s)", "Speaker", "Text"], label="Diarized Transcript", wrap=True, ) run_btn.click( fn=run_pipeline, inputs=[audio_input, json_input], outputs=[output_table, status_box], ) if __name__ == "__main__": demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)