sing_detect/sidecar/detect.py

#!/usr/bin/env python3
"""
Stream Studio — singing-detection sidecar.

This is the ONLY Python in the project. It exists solely because the ML model
(inaSpeechSegmenter, which classifies singing voice as "music") is TensorFlow-
based and has no native Rust/JS equivalent. Everything else — UI, ffmpeg
concat, audio extraction, file exports — is handled by the Tauri/Rust host.

Contract (line-delimited JSON on stdout):
    {"type": "progress", "message": "..."}
    {"type": "result", "segments": [{"index","start","end","duration"}...],
                        "stats": {"music": 123.4, "speech": 456.7, ...}}
    {"type": "error",   "message": "..."}

The host extracts a 16 kHz mono WAV and passes it via --wav, so this script
never touches ffmpeg.

Run with Python 3.10 (TensorFlow does not support 3.12+):
    py -3.10 detect.py --wav audio.wav --gap 10 --min-duration 20 --padding 2
"""

import argparse
import json
import sys
import time


def emit(obj: dict) -> None:
    """Write one JSON record per line and flush immediately."""
    sys.stdout.write(json.dumps(obj, ensure_ascii=False) + "\n")
    sys.stdout.flush()


def progress(msg: str) -> None:
    emit({"type": "progress", "message": msg})


def merge_singing_segments(segments, max_gap, min_duration, padding):
    """Filter for 'music' segments and merge nearby ones into songs."""
    music = sorted((s, e) for label, s, e in segments if label == "music")
    if not music:
        return []

    merged = []
    cur_start, cur_end = music[0]
    for s, e in music[1:]:
        if s - cur_end <= max_gap:
            cur_end = max(cur_end, e)
        else:
            merged.append((cur_start, cur_end))
            cur_start, cur_end = s, e
    merged.append((cur_start, cur_end))

    results = []
    idx = 1
    for s, e in merged:
        if (e - s) >= min_duration:
            ps = max(0.0, s - padding)
            pe = e + padding
            results.append({
                "index": idx,
                "start": round(ps, 3),
                "end": round(pe, 3),
                "duration": round(pe - ps, 3),
            })
            idx += 1
    return results


def _model_present() -> bool:
    """Best-effort check for cached model weights, without downloading."""
    try:
        import glob
        import os

        import inaSpeechSegmenter

        roots = [
            os.path.dirname(inaSpeechSegmenter.__file__),
            os.path.join(os.path.expanduser("~"), ".keras"),
        ]
        for root in roots:
            for ext in ("*.hdf5", "*.h5", "*.keras", "*.pb"):
                if glob.glob(os.path.join(root, "**", ext), recursive=True):
                    return True
        return False
    except Exception:
        return False


def run_check() -> int:
    """Report whether the ML packages (and weights) are available."""
    record = {"type": "check", "packages": False, "model": False, "detail": ""}
    try:
        import inaSpeechSegmenter  # noqa: F401
        import tensorflow  # noqa: F401
        record["packages"] = True
    except Exception as exc:  # noqa: BLE001
        record["detail"] = f"Packages not installed ({type(exc).__name__})."
        emit(record)
        return 0
    record["model"] = _model_present()
    record["detail"] = (
        "Ready — packages and model weights are installed."
        if record["model"]
        else "Packages installed; model weights will download on first run."
    )
    emit(record)
    return 0


def run_warmup() -> int:
    """Instantiate the model once, forcing the weight download if needed."""
    try:
        progress("Loading detection model (downloading weights if needed)…")
        from inaSpeechSegmenter import Segmenter

        Segmenter(vad_engine="smn", detect_gender=False)
        emit({"type": "done", "message": "Model ready."})
        return 0
    except Exception as exc:  # noqa: BLE001
        emit({"type": "error", "message": f"{type(exc).__name__}: {exc}"})
        return 1


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--wav")
    ap.add_argument("--gap", type=float, default=10.0)
    ap.add_argument("--min-duration", type=float, default=20.0)
    ap.add_argument("--padding", type=float, default=2.0)
    ap.add_argument("--check", action="store_true", help="probe environment, then exit")
    ap.add_argument("--warmup", action="store_true", help="download/verify weights, then exit")
    args = ap.parse_args()

    if args.check:
        return run_check()
    if args.warmup:
        return run_warmup()
    if not args.wav:
        emit({"type": "error", "message": "--wav is required for detection."})
        return 1

    try:
        progress("Loading detection model (first run downloads weights)…")
        from inaSpeechSegmenter import Segmenter
    except ImportError:
        emit({
            "type": "error",
            "message": (
                "inaSpeechSegmenter is not installed for this Python.\n"
                "Install it with:\n"
                "  py -3.10 -m pip install inaSpeechSegmenter tensorflow"
            ),
        })
        return 1

    try:
        t0 = time.time()
        seg = Segmenter(vad_engine="smn", detect_gender=False)
        progress("Analysing audio — this can take a while on long streams…")
        raw = seg(args.wav)
        progress(f"Segmentation done in {time.time() - t0:.0f}s ({len(raw)} raw segments).")

        stats = {}
        for label, s, e in raw:
            stats[label] = round(stats.get(label, 0.0) + (e - s), 1)

        songs = merge_singing_segments(raw, args.gap, args.min_duration, args.padding)
        progress(f"Found {len(songs)} singing segment(s) after merge + filter.")
        emit({"type": "result", "segments": songs, "stats": stats})
        return 0
    except Exception as exc:  # noqa: BLE001 — surface any model/runtime error to the host
        emit({"type": "error", "message": f"{type(exc).__name__}: {exc}"})
        return 1


if __name__ == "__main__":
    sys.exit(main())