vidgod/bin/vg-index

#!/bin/sh
"exec" "`dirname $0`/../venvs/index/bin/python" "$0" "$@"
"""vg-index — build a searchable clip library: shot detection + transcript + thumbnails.

Usage: vg-index PATH [PATH...] [--no-whisper] [--force] [--min-shot 0.5]

For every video found under the given paths:
  1. PySceneDetect splits it into shots
  2. a thumbnail is saved for each shot        -> library/thumbs/
  3. Whisper (mlx) transcribes the audio       -> timestamped segments
  4. everything lands in library/clips.sqlite (FTS5 full-text over dialogue)

Search it with vg-find. Re-running skips unchanged files. Heartbeat: ~/.jobs/vidgod-index.status
"""
import argparse, hashlib, os, sqlite3, subprocess, sys, time
from datetime import datetime
from pathlib import Path

os.environ["PATH"] = "/opt/homebrew/bin:" + os.environ.get("PATH", "")

VG = Path(__file__).resolve().parent.parent
LIB = VG / "library"
DB = LIB / "clips.sqlite"
THUMBS = LIB / "thumbs"
WHISPER_MODEL = "mlx-community/whisper-large-v3-turbo"
VIDEO_EXTS = {".mov", ".mp4", ".mkv", ".avi", ".mpg", ".mpeg", ".wmv", ".webm", ".m4v", ".flv", ".mts", ".m2ts", ".vob"}


def hb(msg, note=""):
    j = Path.home() / ".jobs"
    j.mkdir(exist_ok=True)
    (j / "vidgod-index.status").write_text(f"{datetime.now():%Y-%m-%d %H:%M:%S} | {msg} | {note}\n")


def open_db():
    LIB.mkdir(parents=True, exist_ok=True)
    THUMBS.mkdir(parents=True, exist_ok=True)
    db = sqlite3.connect(DB, timeout=60)
    try:  # WAL is persistent once set; the switch itself needs an exclusive
        db.execute("PRAGMA journal_mode=WAL")  # lock, so tolerate losing the race
    except sqlite3.OperationalError:
        pass
    db.executescript("""
    CREATE TABLE IF NOT EXISTS videos(
      id INTEGER PRIMARY KEY, path TEXT UNIQUE, mtime REAL, duration REAL,
      fps REAL, width INT, height INT, has_transcript INT DEFAULT 0, indexed_at TEXT);
    CREATE TABLE IF NOT EXISTS shots(
      id INTEGER PRIMARY KEY, video_id INT, idx INT, start_s REAL, end_s REAL, thumb TEXT);
    CREATE TABLE IF NOT EXISTS segs(
      id INTEGER PRIMARY KEY, video_id INT, start_s REAL, end_s REAL, text TEXT);
    CREATE VIRTUAL TABLE IF NOT EXISTS segs_fts USING fts5(text, content='segs', content_rowid='id');
    CREATE INDEX IF NOT EXISTS idx_shots_video ON shots(video_id);
    CREATE INDEX IF NOT EXISTS idx_segs_video ON segs(video_id);
    """)
    return db


def probe(path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height,avg_frame_rate:format=duration",
         "-of", "csv=p=0", str(path)],
        capture_output=True, text=True).stdout.strip().splitlines()
    w = h = 0; fps = 25.0; dur = 0.0
    for line in out:
        parts = line.split(",")
        if len(parts) >= 3:
            w, h = int(parts[0]), int(parts[1])
            try:
                num, den = parts[2].split("/")
                fps = float(num) / float(den) if float(den) else 25.0
            except ValueError:
                pass
        elif len(parts) == 1 and parts[0]:
            try: dur = float(parts[0])
            except ValueError: pass
    return dur, fps, w, h


def detect_shots(path, min_shot):
    from scenedetect import detect, ContentDetector
    scenes = detect(str(path), ContentDetector(min_scene_len=int(min_shot * 25)))
    if not scenes:
        return None  # single-shot video; caller uses full duration
    return [(s.seconds, e.seconds) for s, e in scenes]


def save_thumb(path, t, dest):
    subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                    "-ss", f"{t:.3f}", "-i", str(path),
                    "-frames:v", "1", "-vf", "scale=480:-2", str(dest)],
                   capture_output=True)
    return dest.exists()


def transcribe(path):
    import mlx_whisper
    res = mlx_whisper.transcribe(str(path), path_or_hf_repo=WHISPER_MODEL)
    return [(s["start"], s["end"], s["text"].strip()) for s in res.get("segments", []) if s["text"].strip()]


def index_file(db, path, args, pos, total):
    st = path.stat()
    row = db.execute("SELECT id, mtime FROM videos WHERE path=?", (str(path),)).fetchone()
    if row and abs(row[1] - st.st_mtime) < 1 and not args.force:
        return "skip"
    if row:  # stale entry — wipe and redo
        vid = row[0]
        db.execute("DELETE FROM segs_fts WHERE rowid IN (SELECT id FROM segs WHERE video_id=?)", (vid,))
        db.execute("DELETE FROM segs WHERE video_id=?", (vid,))
        db.execute("DELETE FROM shots WHERE video_id=?", (vid,))
        db.execute("DELETE FROM videos WHERE id=?", (vid,))

    hb(f"{pos}/{total}", f"probing {path.name}")
    dur, fps, w, h = probe(path)
    if dur <= 0:
        return "unreadable"
    cur = db.execute(
        "INSERT INTO videos(path, mtime, duration, fps, width, height, indexed_at) VALUES(?,?,?,?,?,?,?)",
        (str(path), st.st_mtime, dur, fps, w, h, datetime.now().isoformat(timespec="seconds")))
    vid = cur.lastrowid

    hb(f"{pos}/{total}", f"shots: {path.name}")
    shots = detect_shots(path, args.min_shot) or [(0.0, dur)]
    tag = hashlib.md5(str(path).encode()).hexdigest()[:10]
    for i, (a, b) in enumerate(shots):
        thumb = THUMBS / f"{tag}_{i:04d}.jpg"
        save_thumb(path, a + (b - a) / 2, thumb)
        db.execute("INSERT INTO shots(video_id, idx, start_s, end_s, thumb) VALUES(?,?,?,?,?)",
                   (vid, i, a, b, str(thumb)))

    if not args.no_whisper:
        hb(f"{pos}/{total}", f"whisper: {path.name}")
        try:
            for a, b, text in transcribe(path):
                cur = db.execute("INSERT INTO segs(video_id, start_s, end_s, text) VALUES(?,?,?,?)",
                                 (vid, a, b, text))
                db.execute("INSERT INTO segs_fts(rowid, text) VALUES(?,?)", (cur.lastrowid, text))
            db.execute("UPDATE videos SET has_transcript=1 WHERE id=?", (vid,))
        except Exception as e:  # no audio track etc — keep the shots
            print(f"  whisper failed on {path.name}: {e}", file=sys.stderr)

    db.commit()
    return f"{len(shots)} shots"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("paths", nargs="+")
    ap.add_argument("--no-whisper", action="store_true")
    ap.add_argument("--force", action="store_true", help="re-index even if unchanged")
    ap.add_argument("--min-shot", type=float, default=0.5, help="minimum shot length, seconds")
    args = ap.parse_args()

    files = []
    for p in args.paths:
        p = Path(p).expanduser().resolve()
        if p.is_dir():
            files += sorted(x for x in p.rglob("*") if x.suffix.lower() in VIDEO_EXTS and x.is_file())
        elif p.is_file():
            files.append(p)
    if not files:
        sys.exit("no video files found")

    db = open_db()
    t0 = time.time()
    counts = {}
    for i, f in enumerate(files, 1):
        try:
            r = index_file(db, f, args, i, len(files))
        except KeyboardInterrupt:
            raise
        except Exception as e:
            r = "error"
            print(f"ERROR indexing {f}: {e}", file=sys.stderr)
        counts[r] = counts.get(r, 0) + 1
        print(f"[{i}/{len(files)}] {f.name}: {r}")
    hb(f"DONE {len(files)}/{len(files)}", str(counts))
    n = db.execute("SELECT COUNT(*) FROM shots").fetchone()[0]
    s = db.execute("SELECT COUNT(*) FROM segs").fetchone()[0]
    print(f"\nlibrary now: {n} shots, {s} dialogue segments ({time.time()-t0:.0f}s). Search with vg-find.")


if __name__ == "__main__":
    main()