app.py

"""
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("<if", header)
            raw = f.read(n_samples * 4)
        samples = np.frombuffer(raw, dtype="<f4").astype(np.float32)
        return samples, PCM_SAMPLE_RATE
    else:
        audio = AudioSegment.from_file(audio_path).set_channels(1).set_frame_rate(PCM_SAMPLE_RATE)
        samples = np.array(audio.get_array_of_samples()).astype(np.float32) / 32768.0
        return samples, audio.frame_rate


def get_embeddings_batched(chunks: list, batch_size: int = 16) -> 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)