From 87e571c4def9a86c1049a0037ef72c344a0b9f8c Mon Sep 17 00:00:00 2001 From: CosmoWorker Date: Sat, 9 May 2026 23:26:44 +0530 Subject: [PATCH 1/2] added single pipeline flow for detection of events via hf models & srt muxing --- main.py | 338 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 338 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..6b5a0bf --- /dev/null +++ b/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: