mic_diarize_plot.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "mlx-audio",
#     "numpy",
#     "pyqtgraph",
#     "PyQt6",
#     "sounddevice",
# ]
# ///
"""Real-time speaker diarization with live PyQtGraph visualization."""

import argparse
import sys
import time
from collections import defaultdict
from dataclasses import dataclass, field
from queue import Empty, Queue

import numpy as np
import pyqtgraph as pg
import sounddevice as sd
from PyQt6.QtCore import QTimer
from PyQt6.QtGui import QColor, QFont
from PyQt6.QtWidgets import QApplication, QLabel, QVBoxLayout, QWidget

SPEAKER_COLORS = [
    "#58a6ff",  # blue
    "#f78166",  # red/orange
    "#7ee787",  # green
    "#d2a8ff",  # purple
]
BG_COLOR = "#0d1117"
SURFACE_COLOR = "#161b22"
BORDER_COLOR = "#30363d"
TEXT_COLOR = "#c9d1d9"
DIM_COLOR = "#484f58"
PLAYHEAD_COLOR = "#f0883e"
VIEW_WINDOW = 30.0


def speaker_label(spk: int) -> str:
    return f"Speaker {chr(65 + spk)}"


def format_time(seconds: float) -> str:
    m, s = divmod(seconds, 60)
    return f"{int(m):02d}:{s:05.2f}"


@dataclass
class Segment:
    speaker: int
    start: float
    end: float


@dataclass
class SessionState:
    segments: list[Segment] = field(default_factory=list)
    speaker_time: dict = field(default_factory=lambda: defaultdict(float))

    def add_segments(self, new_segs):
        for seg in new_segs:
            s = Segment(speaker=seg.speaker, start=seg.start, end=seg.end)
            # Merge with last segment if same speaker and contiguous
            if (
                self.segments
                and self.segments[-1].speaker == s.speaker
                and abs(self.segments[-1].end - s.start) < 0.01
            ):
                self.segments[-1].end = s.end
            else:
                self.segments.append(s)

    def prune(self, cutoff: float):
        """Drop segments that end before cutoff, clip those that straddle it."""
        pruned = []
        for seg in self.segments:
            if seg.end <= cutoff:
                continue
            if seg.start < cutoff:
                seg.start = cutoff
            pruned.append(seg)
        self.segments = pruned

    def compute_speaker_time(self) -> dict[int, float]:
        totals: dict[int, float] = defaultdict(float)
        for seg in self.segments:
            totals[seg.speaker] += seg.end - seg.start
        return totals


class DiarizationWindow(QWidget):
    def __init__(self, model, args):
        super().__init__()
        self.model = model
        self.args = args
        self.state = model.init_streaming_state()
        self.session = SessionState()
        self.t0 = time.time()
        self.audio_queue: Queue = Queue()
        self.last_chunk_end: float = 0.0  # wall time of last processed chunk's end

        self._build_ui()
        self._start_audio()
        self._start_timer()

    def _build_ui(self):
        self.setWindowTitle("Live Speaker Diarization")
        self.setStyleSheet(
            f"background-color: {BG_COLOR}; color: {TEXT_COLOR};"
        )
        self.resize(1100, 450)

        layout = QVBoxLayout(self)
        layout.setContentsMargins(16, 12, 16, 12)
        layout.setSpacing(8)

        # -- Header --
        header_font = QFont("SF Mono, Menlo, monospace", 13)
        header_font.setBold(True)
        self.header_label = QLabel()
        self.header_label.setFont(header_font)
        self.header_label.setStyleSheet(f"color: {TEXT_COLOR}; padding: 4px 0;")
        layout.addWidget(self.header_label)

        # -- Speaker timeline --
        self.timeline_plot = pg.PlotWidget()
        self.timeline_plot.setBackground(SURFACE_COLOR)
        self.timeline_plot.setTitle("Speaker Timeline", color=TEXT_COLOR, size="11pt")
        self.timeline_plot.setLabel("bottom", "Time (s)", color=DIM_COLOR)
        self.timeline_plot.showGrid(x=True, y=False, alpha=0.15)
        self.timeline_plot.setMouseEnabled(x=True, y=False)
        self.timeline_plot.setYRange(-0.5, 3.5)

        # Y-axis speaker labels
        yticks = [(i, speaker_label(i)) for i in range(4)]
        self.timeline_plot.getAxis("left").setTicks([yticks])
        self.timeline_plot.getAxis("left").setStyle(tickTextOffset=8)
        self.timeline_plot.getAxis("left").setTextPen(pg.mkPen(TEXT_COLOR))

        # Processing zone — gray region for buffered but not-yet-diarized audio
        self.processing_region = pg.LinearRegionItem(
            values=(0, 0), orientation="vertical", movable=False,
            brush=pg.mkBrush(255, 255, 255, 25),
            pen=pg.mkPen(None),
        )
        self.processing_region.setZValue(-1)
        self.timeline_plot.addItem(self.processing_region)

        # Playhead line on timeline
        self.playhead = pg.InfiniteLine(
            pos=0, angle=90,
            pen=pg.mkPen(PLAYHEAD_COLOR, width=2, style=pg.QtCore.Qt.PenStyle.DashLine),
        )
        self.timeline_plot.addItem(self.playhead)

        layout.addWidget(self.timeline_plot, stretch=1)

        # -- Stats bar --
        stats_font = QFont("SF Mono, Menlo, monospace", 11)
        self.stats_label = QLabel()
        self.stats_label.setFont(stats_font)
        self.stats_label.setStyleSheet(
            f"color: {DIM_COLOR}; padding: 4px 0;"
        )
        layout.addWidget(self.stats_label)

        self._update_header(0.0)

    def _start_audio(self):
        chunk_samples = int(self.args.chunk_duration * self.args.sample_rate)

        def callback(indata, frames, time_info, status):
            if status:
                print(f"[audio: {status}]", file=sys.stderr)
            self.audio_queue.put(indata[:, 0].copy())

        self.stream = sd.InputStream(
            samplerate=self.args.sample_rate,
            channels=1,
            dtype="float32",
            blocksize=chunk_samples,
            device=self.args.device,
            callback=callback,
        )
        self.stream.start()

    def _start_timer(self):
        self.timer = QTimer()
        self.timer.timeout.connect(self._tick)
        self.timer.start(50)  # 20 fps

    def _tick(self):
        elapsed = time.time() - self.t0
        # Drain audio queue
        while True:
            try:
                chunk = self.audio_queue.get_nowait()
            except Empty:
                break

            result, self.state = self.model.feed(
                chunk,
                self.state,
                sample_rate=self.args.sample_rate,
                threshold=self.args.threshold,
            )
            self.session.add_segments(result.segments)
            self.last_chunk_end = elapsed

        # Prune segments outside the view window
        if elapsed > VIEW_WINDOW:
            self.session.prune(elapsed - VIEW_WINDOW)

        # Processing zone: from last diarized point to current time
        self.processing_region.setRegion((self.last_chunk_end, elapsed))

        # Update timeline bars
        self._draw_timeline(elapsed)
        self._update_header(elapsed)
        self._update_stats(elapsed)

    def _draw_timeline(self, elapsed: float):
        # Remove old bar items
        for item in list(self.timeline_plot.items()):
            if isinstance(item, pg.BarGraphItem):
                self.timeline_plot.removeItem(item)

        if not self.session.segments:
            self.playhead.setValue(elapsed)
            return

        # Group segments by speaker for batch drawing
        by_speaker: dict[int, list[Segment]] = defaultdict(list)
        for seg in self.session.segments:
            by_speaker[seg.speaker].append(seg)

        for spk, segs in by_speaker.items():
            x0 = np.array([s.start for s in segs])
            widths = np.array([s.end - s.start for s in segs])
            y0 = np.full(len(segs), spk - 0.35)
            heights = np.full(len(segs), 0.7)
            color = QColor(SPEAKER_COLORS[spk % len(SPEAKER_COLORS)])
            color.setAlpha(200)
            bars = pg.BarGraphItem(
                x0=x0, width=widths, y0=y0, height=heights,
                brush=color,
                pen=pg.mkPen(SPEAKER_COLORS[spk % len(SPEAKER_COLORS)], width=1),
            )
            self.timeline_plot.addItem(bars)

        # Auto-scroll timeline: show last 30s or full range
        view_window = 30.0
        if elapsed > view_window:
            self.timeline_plot.setXRange(elapsed - view_window, elapsed + 1)
        else:
            self.timeline_plot.setXRange(0, max(elapsed + 1, 5))

        self.playhead.setValue(elapsed)

    def _update_header(self, elapsed: float):
        self.header_label.setText(
            f"  Recording: {format_time(elapsed)}    "
            f"Segments: {len(self.session.segments)}    "
            f"Speakers: {len(set(s.speaker for s in self.session.segments))}"
        )

    def _update_stats(self, elapsed: float):
        speaker_time = self.session.compute_speaker_time()
        window = min(elapsed, VIEW_WINDOW)
        parts = []
        for spk in sorted(speaker_time):
            t = speaker_time[spk]
            pct = (t / window * 100) if window > 0 else 0
            label = speaker_label(spk)
            parts.append(f"{label}: {t:.1f}s ({pct:.0f}%)")
        self.stats_label.setText("  " + "    ".join(parts) if parts else "  Listening...")

    def closeEvent(self, event):
        self.timer.stop()
        self.stream.stop()
        self.stream.close()
        event.accept()


def main():
    parser = argparse.ArgumentParser(description="Real-time diarization with live plot")
    parser.add_argument(
        "--model",
        default="mlx-community/diar_streaming_sortformer_4spk-v2.1-fp32",
        help="Model path or HuggingFace repo ID",
    )
    parser.add_argument(
        "--chunk-duration", type=float, default=3.0,
        help="Audio chunk duration in seconds (default: 3.0)",
    )
    parser.add_argument(
        "--sample-rate", type=int, default=16000,
        help="Mic sample rate (default: 16000)",
    )
    parser.add_argument(
        "--threshold", type=float, default=0.5,
        help="Speaker activity threshold (default: 0.5)",
    )
    parser.add_argument(
        "--device", type=int, default=None,
        help="Audio input device index",
    )
    parser.add_argument(
        "--list-devices", action="store_true",
        help="List audio devices and exit",
    )
    args = parser.parse_args()

    if args.list_devices:
        print(sd.query_devices())
        return

    print("Loading model...")
    from mlx_audio.vad import load
    model = load(args.model)
    print("Ready.\n")

    app = QApplication(sys.argv)
    win = DiarizationWindow(model, args)
    win.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()