Initial Implementation of CC generation via HF models

.env.example

@@ -0,0 +1 @@
+HF_TOKEN=

.gitignore

@@ -0,0 +1,5 @@
+.venv
+.env
+*.mp4
+*.srt
+__pycache__

.python-version

@@ -0,0 +1 @@
+3.12

main.py

@@ -0,0 +1,338 @@
+#!/usr/bin/env python3
+"""
+Intelligent Closed Caption (CC) Suggestion Tool — Prototype v2
+Detects non-speech audio events in video and generates .srt subtitle files.
+"""
+
+import os
+# import sys
+import argparse
+import subprocess
+import tempfile
+from dataclasses import dataclass
+from datetime import timedelta
+from pathlib import Path
+from typing import Optional
+from dotenv import load_dotenv
+
+import cv2
+import numpy as np
+import librosa
+# import soundfile as sf
+import srt
+from transformers import pipeline as hf_pipeline
+
+# CONFIG
+load_dotenv()
+HF_TOKEN = os.getenv("HF_TOKEN")
+MODEL_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"
+SAMPLE_RATE = 16_000
+WINDOW_SEC = 2.0
+HOP_SEC = 0.5 
+
+
+SKIP_TERMS = [
+    "Speech",
+    "Narration",
+    "Male speech",
+    "Female speech",
+    "Child speech",
+    "Conversation",
+    "Silence",
+    "White noise",
+    "Pink noise",
+    "Static",
+    "Background noise",
+    "Noise",
+]
+
+# Minimum confidence for a non-speech label to count in a window
+# Lower than v1 so background events in talk-heavy videos register
+AUDIO_CONF_THRESHOLD = 0.12
+
+# Accept merged event only if its averaged window confidence clears this
+EVENT_CONF_THRESHOLD = 0.20
+
+# Final weighted score threshold
+FINAL_SCORE_THRESHOLD = 0.30
+
+# Weights
+W_AUDIO = 0.70
+W_VISUAL = 0.30
+
+AUTO_TRIGGER_CLASSES = {"Explosion", "Gunshot", "Screaming", "Alarm", "Siren"}
+AUTO_TRIGGER_THRESHOLD = 0.80
+
+# Max gap (seconds) between same-label windows before starting a new event
+MERGE_GAP_SEC = 1.0
+
+# Minimum event duration to write to SRT
+MIN_EVENT_SEC = 0.5
+
+
+
+@dataclass
+class AudioEvent:
+    label: str
+    confidence: float
+    start_time: float
+    end_time: float
+    visual_score: float = 0.0
+    final_score: float = 0.0
+    accepted: bool = False
+
+    @property
+    def duration(self) -> float:
+        return self.end_time - self.start_time
+
+
+def extract_audio(video_path: str, wav_path: str) -> str:
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        video_path,
+        "-ar",
+        str(SAMPLE_RATE),
+        "-ac",
+        "1",
+        "-vn",
+        wav_path,
+    ]
+    result = subprocess.run(cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"FFmpeg failed:\n{result.stderr}")
+    return wav_path
+
+
+def _is_speech_or_noise(label: str) -> bool:
+    return any(t.lower() in label.lower() for t in SKIP_TERMS)
+
+
+def _best_non_speech(results: list[dict]) -> Optional[tuple[str, float]]:
+    """Return highest-scoring non-speech label above AUDIO_CONF_THRESHOLD."""
+    for r in results:
+        label, score = r["label"], r["score"]
+        if _is_speech_or_noise(label):
+            continue
+        if score >= AUDIO_CONF_THRESHOLD:
+            return label, score
+    return None
+
+
+def detect_audio_events(wav_path: str, debug: bool = False) -> list["AudioEvent"]:
+    print("  Loading AST model …")
+    classifier = hf_pipeline(
+        "audio-classification",
+        model=MODEL_NAME,
+        token=HF_TOKEN,
+        top_k=10,  # larger K → better chance of catching background events
+        device=-1,
+    )
+
+    audio, _ = librosa.load(wav_path, sr=SAMPLE_RATE, mono=True)
+    duration = len(audio) / SAMPLE_RATE
+    print(f"  Duration: {duration:.1f}s | window={WINDOW_SEC}s, hop={HOP_SEC}s")
+
+    window_samples = int(WINDOW_SEC * SAMPLE_RATE)
+    hop_samples = int(HOP_SEC * SAMPLE_RATE)
+
+    raw: list[tuple[float, float, str, float]] = []
+
+    for start in range(0, len(audio) - window_samples + 1, hop_samples):
+        chunk = audio[start : start + window_samples]
+        start_t = round(start / SAMPLE_RATE, 2)
+        end_t = round(start_t + WINDOW_SEC, 2)
+
+        results = classifier({"array": chunk, "sampling_rate": SAMPLE_RATE})
+
+        if debug:
+            top3 = [(r["label"][:35], f"{r['score']:.3f}") for r in results[:3]]
+            print(f"    [{start_t:5.1f}s] " + " | ".join(f"{l} {s}" for l, s in top3))
+
+        hit = _best_non_speech(results)
+        if hit:
+            raw.append((start_t, end_t, hit[0], hit[1]))
+
+    print(f"  Windows with non-speech hits: {len(raw)}")
+    return _merge_events(raw)
+
+
+def _merge_events(raw: list[tuple]) -> list[AudioEvent]:
+    if not raw:
+        return []
+
+    raw.sort(key=lambda x: x[0])
+    events: list[AudioEvent] = []
+
+    cur_start, cur_end, cur_label, _ = raw[0]
+    conf_pool = [raw[0][3]]
+
+    for start, end, label, conf in raw[1:]:
+        if label == cur_label and start <= cur_end + MERGE_GAP_SEC:
+            cur_end = max(cur_end, end)
+            conf_pool.append(conf)
+        else:
+            avg_conf = float(np.mean(conf_pool))
+            if avg_conf >= EVENT_CONF_THRESHOLD:
+                events.append(
+                    AudioEvent(
+                        label=cur_label,
+                        confidence=round(avg_conf, 4),
+                        start_time=cur_start,
+                        end_time=cur_end,
+                    )
+                )
+            cur_start, cur_end, cur_label = start, end, label
+            conf_pool = [conf]
+
+    avg_conf = float(np.mean(conf_pool))
+    if avg_conf >= EVENT_CONF_THRESHOLD:
+        events.append(
+            AudioEvent(
+                label=cur_label,
+                confidence=round(avg_conf, 4),
+                start_time=cur_start,
+                end_time=cur_end,
+            )
+        )
+
+    return [e for e in events if e.duration >= MIN_EVENT_SEC]
+
+def _read_frame(cap: cv2.VideoCapture, t_sec: float) -> Optional[np.ndarray]:
+    cap.set(cv2.CAP_PROP_POS_MSEC, t_sec * 1000)
+    ret, frame = cap.read()
+    return frame if ret else None
+
+
+def validate_visually(video_path: str, event: AudioEvent) -> float:
+    cap = cv2.VideoCapture(video_path)
+    t_before = max(0.0, event.start_time - 0.3)
+    t_after = event.start_time + 0.5
+    f1 = _read_frame(cap, t_before)
+    f2 = _read_frame(cap, t_after)
+    cap.release()
+
+    if f1 is None or f2 is None:
+        return 0.3
+
+    g1 = cv2.cvtColor(f1, cv2.COLOR_BGR2GRAY)
+    g2 = cv2.cvtColor(f2, cv2.COLOR_BGR2GRAY)
+    diff = cv2.absdiff(g1, g2)
+    motion_raw = float(diff.mean()) / 255.0
+    visual_score = min(1.0, motion_raw * 5.0)
+    return round(visual_score, 4)
+
+
+
+def decide(event: AudioEvent, visual_score: float) -> tuple[float, bool]:
+    if (
+        event.label in AUTO_TRIGGER_CLASSES
+        and event.confidence >= AUTO_TRIGGER_THRESHOLD
+    ):
+        return round(event.confidence, 4), True
+    final = round((event.confidence * W_AUDIO) + (visual_score * W_VISUAL), 4)
+    accepted = final >= FINAL_SCORE_THRESHOLD
+    return final, accepted
+
+
+def _fmt_label(raw_label: str) -> str:
+    label = raw_label.lower().strip()
+    for sep in [",", "("]:
+        label = label.split(sep)[0].strip()
+    return f"[{label}]"
+
+
+def generate_srt(events: list[AudioEvent], output_path: str) -> str:
+    subtitles = [
+        srt.Subtitle(
+            index=i + 1,
+            start=timedelta(seconds=e.start_time),
+            end=timedelta(seconds=e.end_time),
+            content=_fmt_label(e.label),
+        )
+        for i, e in enumerate(events)
+        if e.accepted
+    ]
+    content = srt.compose(subtitles)
+    Path(output_path).write_text(content, encoding="utf-8")
+    return content
+
+
+# main pipeline
+def run(video_path: str, output_srt: str, debug: bool = False) -> None:
+    sep = "─" * 64
+    print(f"\n{sep}\nCC Tool v2  ·  {video_path}\n{sep}")
+
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+        wav_path = f.name
+
+    try:
+        print("\n[Phase 1] Extracting audio …")
+        extract_audio(video_path, wav_path)
+
+        print("[Phase 1] Detecting audio events …")
+        events = detect_audio_events(wav_path, debug=debug)
+        print(f"  → {len(events)} merged candidate event(s)")
+
+        if not events:
+            print("\n  ⚠  No non-speech events detected above threshold.")
+            print("     Run with --debug to inspect raw per-window predictions.")
+
+        print("\n[Phase 2+3] Visual validation & scoring …")
+        hdr = f"  {'STATUS':<12} {'LABEL':<28} {'AUD':>5} {'VIS':>5} {'FIN':>5}  SPAN"
+        print(hdr)
+        print("  " + "─" * (len(hdr) - 2))
+
+        for event in events:
+            vis = validate_visually(video_path, event)
+            final, accept = decide(event, vis)
+            event.visual_score = vis
+            event.final_score = final
+            event.accepted = accept
+
+            tick = "✓ ACCEPT" if accept else "✗ reject"
+            label = event.label[:26]
+            print(
+                f"  {tick:<12} {label:<28} "
+                f"{event.confidence:>5.2f} {vis:>5.2f} {final:>5.2f}  "
+                f"[{event.start_time:.1f}s–{event.end_time:.1f}s]"
+            )
+
+        print(f"\n[Phase 4] Writing SRT → {output_srt}")
+        srt_text = generate_srt(events, output_srt)
+
+        accepted = [e for e in events if e.accepted]
+        print(f"\n{len(accepted)}/{len(events)} events written to '{output_srt}'")
+
+        if srt_text.strip():
+            lines = srt_text.splitlines()
+            preview = "\n".join(lines[: min(30, len(lines))])
+            print(f"\n── SRT Preview ──────────────────\n{preview}")
+
+    finally:
+        if os.path.exists(wav_path):
+            os.unlink(wav_path)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate non-speech closed captions for a video file."
+    )
+    parser.add_argument("video", help="Input video file path")
+    parser.add_argument(
+        "output", nargs="?", help="Output .srt path (default: <video>_cc.srt)"
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Print top-3 model predictions for every window",
+    )
+    args = parser.parse_args()
+
+    output = args.output or (Path(args.video).stem + "_cc.srt")
+    run(args.video, output, debug=args.debug)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "intelligent-cc-generation-tmp"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "librosa>=0.11.0",
+    "numpy>=2.4.4",
+    "opencv-python>=4.13.0.92",
+    "python-dotenv>=1.2.2",
+    "soundfile>=0.13.1",
+    "srt>=3.5.3",
+    "torch>=2.11.0",
+    "torchaudio>=2.11.0",
+    "transformers>=5.8.0",
+]