|
| 1 | +""" |
| 2 | +Pipelined per-video processing: transcribe → extract frames → align. |
| 3 | +
|
| 4 | +Uses a producer-consumer pattern with two threads so that while video N+1 |
| 5 | +is being transcribed (GPU), video N can be frame-extracted and aligned |
| 6 | +concurrently. This overlaps the stages and reduces total wall-clock time |
| 7 | +compared to the sequential approach (transcribe ALL → frame-extract ALL → |
| 8 | +align ALL). |
| 9 | +
|
| 10 | +Usage: |
| 11 | + python pipeline_worker.py --course 85397 [--force] [--path ~/AutoNote] |
| 12 | +
|
| 13 | +Replaces the three separate calls to: |
| 14 | + extract_caption.py / frame_extractor.py / semantic_alignment.py |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +import argparse |
| 20 | +import json |
| 21 | +import os |
| 22 | +import subprocess |
| 23 | +import sys |
| 24 | +import threading |
| 25 | +from pathlib import Path |
| 26 | +from queue import Queue |
| 27 | + |
| 28 | +PROJECT_DIR = Path(__file__).parent |
| 29 | +_AUTO_NOTE_DIR = Path.home() / ".auto_note" |
| 30 | +if os.environ.get("AUTONOTE_DATA_DIR"): |
| 31 | + DATA_DIR = Path(os.environ["AUTONOTE_DATA_DIR"]) |
| 32 | +elif getattr(sys, "frozen", False) or PROJECT_DIR == _AUTO_NOTE_DIR / "scripts": |
| 33 | + DATA_DIR = _AUTO_NOTE_DIR |
| 34 | +else: |
| 35 | + DATA_DIR = PROJECT_DIR |
| 36 | + |
| 37 | +try: |
| 38 | + _cfg_file = DATA_DIR / "config.json" |
| 39 | + _cfg: dict = json.loads(_cfg_file.read_text(encoding="utf-8")) if _cfg_file.exists() else {} |
| 40 | +except Exception: |
| 41 | + _cfg = {} |
| 42 | +_out = _cfg.get("OUTPUT_DIR", "").strip() |
| 43 | +COURSE_DATA_DIR = Path(_out) if _out else Path.home() / "AutoNote" |
| 44 | + |
| 45 | +MANIFEST_FILE = DATA_DIR / "manifest.json" |
| 46 | + |
| 47 | +# Detect scripts location |
| 48 | +SCRIPTS_DIR = _AUTO_NOTE_DIR / "scripts" |
| 49 | +if not SCRIPTS_DIR.exists(): |
| 50 | + SCRIPTS_DIR = PROJECT_DIR |
| 51 | + |
| 52 | +PYTHON = sys.executable |
| 53 | + |
| 54 | +# Prevent console flashing on Windows |
| 55 | +_CREATIONFLAGS = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0 |
| 56 | + |
| 57 | + |
| 58 | +def _script(name: str) -> str: |
| 59 | + p = SCRIPTS_DIR / name |
| 60 | + if p.exists(): |
| 61 | + return str(p) |
| 62 | + return str(PROJECT_DIR / name) |
| 63 | + |
| 64 | + |
| 65 | +def _run(cmd: list[str], label: str) -> bool: |
| 66 | + """Run a subprocess, streaming output to stdout.""" |
| 67 | + print(f"\n{'─' * 60}") |
| 68 | + print(f" {label}") |
| 69 | + print(f"{'─' * 60}", flush=True) |
| 70 | + proc = subprocess.run( |
| 71 | + cmd, |
| 72 | + env={**os.environ, "PYTHONUNBUFFERED": "1"}, |
| 73 | + creationflags=_CREATIONFLAGS, |
| 74 | + ) |
| 75 | + return proc.returncode == 0 |
| 76 | + |
| 77 | + |
| 78 | +def get_videos(course_id: str, base_dir: Path) -> list[dict]: |
| 79 | + """Get list of downloaded videos for a course from manifest.""" |
| 80 | + if not MANIFEST_FILE.exists(): |
| 81 | + return [] |
| 82 | + with open(MANIFEST_FILE, encoding="utf-8") as f: |
| 83 | + manifest = json.load(f) |
| 84 | + |
| 85 | + videos = [] |
| 86 | + for key, entry in manifest.items(): |
| 87 | + if entry.get("status") != "done": |
| 88 | + continue |
| 89 | + vpath = entry.get("path", "") |
| 90 | + if not vpath or not Path(vpath).exists(): |
| 91 | + continue |
| 92 | + try: |
| 93 | + if Path(vpath).parent.parent.name != course_id: |
| 94 | + continue |
| 95 | + except Exception: |
| 96 | + continue |
| 97 | + videos.append({ |
| 98 | + "path": Path(vpath), |
| 99 | + "stem": Path(vpath).stem, |
| 100 | + "stream_tag": entry.get("stream_tag", ""), |
| 101 | + "course_dir": base_dir / course_id, |
| 102 | + }) |
| 103 | + return videos |
| 104 | + |
| 105 | + |
| 106 | +def transcribe_one(video: dict, force: bool) -> bool: |
| 107 | + """Transcribe a single video.""" |
| 108 | + caption = video["course_dir"] / "captions" / f"{video['stem']}.json" |
| 109 | + if not force and caption.exists(): |
| 110 | + print(f" [skip] Already transcribed: {video['stem']}") |
| 111 | + return True |
| 112 | + cmd = [PYTHON, _script("extract_caption.py"), "--video", str(video["path"])] |
| 113 | + if force: |
| 114 | + cmd.append("--force") |
| 115 | + return _run(cmd, f"Transcribe: {video['stem']}") |
| 116 | + |
| 117 | + |
| 118 | +def extract_frames_one(video: dict) -> bool: |
| 119 | + """Extract frames from a single video (if screenshare).""" |
| 120 | + caption = video["course_dir"] / "captions" / f"{video['stem']}.json" |
| 121 | + if not caption.exists(): |
| 122 | + return True # no caption yet, skip frame extraction |
| 123 | + align_file = video["course_dir"] / "alignment" / f"{video['stem']}.json" |
| 124 | + # Skip if already has screenshare alignment |
| 125 | + if align_file.exists(): |
| 126 | + try: |
| 127 | + data = json.loads(align_file.read_text(encoding="utf-8")) |
| 128 | + if data.get("source") == "screenshare": |
| 129 | + print(f" [skip] Frames already extracted: {video['stem']}") |
| 130 | + return True |
| 131 | + except Exception: |
| 132 | + pass |
| 133 | + |
| 134 | + cmd = [PYTHON, _script("frame_extractor.py"), |
| 135 | + "--video", str(video["path"]), |
| 136 | + "--caption", str(caption), |
| 137 | + "--course-dir", str(video["course_dir"])] |
| 138 | + return _run(cmd, f"Extract frames: {video['stem']}") |
| 139 | + |
| 140 | + |
| 141 | +def align_one(video: dict, force: bool) -> bool: |
| 142 | + """Align a single caption to slides.""" |
| 143 | + caption = video["course_dir"] / "captions" / f"{video['stem']}.json" |
| 144 | + if not caption.exists(): |
| 145 | + return True # no caption, skip |
| 146 | + |
| 147 | + # Skip if screenshare alignment exists (frame_extractor handled it) |
| 148 | + align_file = video["course_dir"] / "alignment" / f"{video['stem']}.json" |
| 149 | + if not force and align_file.exists(): |
| 150 | + try: |
| 151 | + data = json.loads(align_file.read_text(encoding="utf-8")) |
| 152 | + if data.get("source") == "screenshare": |
| 153 | + print(f" [skip] Screenshare alignment exists: {video['stem']}") |
| 154 | + return True |
| 155 | + except Exception: |
| 156 | + pass |
| 157 | + |
| 158 | + # For slide-based alignment, use the course-level command |
| 159 | + # (it handles matching captions to slides) |
| 160 | + if not force and align_file.exists(): |
| 161 | + print(f" [skip] Already aligned: {video['stem']}") |
| 162 | + return True |
| 163 | + |
| 164 | + cmd = [PYTHON, _script("semantic_alignment.py"), |
| 165 | + "--caption", str(caption)] |
| 166 | + |
| 167 | + # Find matching slides — use the course-level discovery |
| 168 | + # Just run for the whole course; it will skip already-aligned captions |
| 169 | + cmd = [PYTHON, _script("semantic_alignment.py"), |
| 170 | + "--course", video["course_dir"].name] |
| 171 | + if force: |
| 172 | + cmd.append("--force") |
| 173 | + return _run(cmd, f"Align: {video['stem']}") |
| 174 | + |
| 175 | + |
| 176 | +def pipeline_sequential(videos: list[dict], force: bool) -> None: |
| 177 | + """Simple sequential pipeline: transcribe → frames → align per video.""" |
| 178 | + for i, video in enumerate(videos, 1): |
| 179 | + print(f"\n{'═' * 60}") |
| 180 | + print(f" Video {i}/{len(videos)}: {video['stem']}") |
| 181 | + print(f"{'═' * 60}", flush=True) |
| 182 | + |
| 183 | + transcribe_one(video, force) |
| 184 | + extract_frames_one(video) |
| 185 | + |
| 186 | + # Alignment is best done course-level (batch BGE-M3 matching) |
| 187 | + if videos: |
| 188 | + course_dir = videos[0]["course_dir"] |
| 189 | + cmd = [PYTHON, _script("semantic_alignment.py"), |
| 190 | + "--course", course_dir.name] |
| 191 | + if force: |
| 192 | + cmd.append("--force") |
| 193 | + _run(cmd, "Align all transcripts") |
| 194 | + |
| 195 | + |
| 196 | +def pipeline_threaded(videos: list[dict], force: bool) -> None: |
| 197 | + """Threaded pipeline: transcribe and frame-extract/align overlap. |
| 198 | +
|
| 199 | + Thread 1 (transcriber): transcribes videos one at a time, pushes |
| 200 | + completed videos into a queue. |
| 201 | + Thread 2 (processor): takes transcribed videos and extracts frames. |
| 202 | + After all videos are processed, alignment runs once (batch mode). |
| 203 | + """ |
| 204 | + if not videos: |
| 205 | + return |
| 206 | + |
| 207 | + queue: Queue[dict | None] = Queue() |
| 208 | + errors: list[str] = [] |
| 209 | + |
| 210 | + def transcriber(): |
| 211 | + for i, video in enumerate(videos, 1): |
| 212 | + print(f"\n{'═' * 60}") |
| 213 | + print(f" [{i}/{len(videos)}] Transcribing: {video['stem']}") |
| 214 | + print(f"{'═' * 60}", flush=True) |
| 215 | + ok = transcribe_one(video, force) |
| 216 | + if not ok: |
| 217 | + errors.append(f"Transcribe failed: {video['stem']}") |
| 218 | + queue.put(video) |
| 219 | + queue.put(None) # sentinel |
| 220 | + |
| 221 | + def frame_processor(): |
| 222 | + while True: |
| 223 | + video = queue.get() |
| 224 | + if video is None: |
| 225 | + break |
| 226 | + print(f"\n Processing frames: {video['stem']}", flush=True) |
| 227 | + extract_frames_one(video) |
| 228 | + |
| 229 | + t1 = threading.Thread(target=transcriber, name="transcriber") |
| 230 | + t2 = threading.Thread(target=frame_processor, name="frame-processor") |
| 231 | + |
| 232 | + t1.start() |
| 233 | + t2.start() |
| 234 | + t1.join() |
| 235 | + t2.join() |
| 236 | + |
| 237 | + # Alignment runs after all transcriptions + frame extractions are done |
| 238 | + # (needs BGE-M3 model which conflicts with Whisper GPU usage) |
| 239 | + course_dir = videos[0]["course_dir"] |
| 240 | + cmd = [PYTHON, _script("semantic_alignment.py"), |
| 241 | + "--course", course_dir.name] |
| 242 | + if force: |
| 243 | + cmd.append("--force") |
| 244 | + _run(cmd, "Align all transcripts") |
| 245 | + |
| 246 | + if errors: |
| 247 | + print(f"\n[warn] {len(errors)} error(s):") |
| 248 | + for e in errors: |
| 249 | + print(f" - {e}") |
| 250 | + |
| 251 | + |
| 252 | +def main() -> None: |
| 253 | + parser = argparse.ArgumentParser( |
| 254 | + description="Pipelined per-video processing: transcribe → frames → align") |
| 255 | + parser.add_argument("--course", required=True, help="Course ID") |
| 256 | + parser.add_argument("--path", help="Base output directory") |
| 257 | + parser.add_argument("--force", action="store_true", |
| 258 | + help="Re-process even if output files exist") |
| 259 | + parser.add_argument("--sequential", action="store_true", |
| 260 | + help="Disable threading (debug mode)") |
| 261 | + args = parser.parse_args() |
| 262 | + |
| 263 | + base_dir = Path(args.path) if args.path else COURSE_DATA_DIR |
| 264 | + videos = get_videos(args.course, base_dir) |
| 265 | + |
| 266 | + if not videos: |
| 267 | + print(f"No downloaded videos found for course {args.course}.") |
| 268 | + print("Run 'Download videos' first.") |
| 269 | + return |
| 270 | + |
| 271 | + print(f"Found {len(videos)} video(s) for course {args.course}.") |
| 272 | + |
| 273 | + if args.sequential: |
| 274 | + pipeline_sequential(videos, args.force) |
| 275 | + else: |
| 276 | + pipeline_threaded(videos, args.force) |
| 277 | + |
| 278 | + print(f"\n✓ Pipeline complete: {len(videos)} video(s) processed.") |
| 279 | + |
| 280 | + |
| 281 | +if __name__ == "__main__": |
| 282 | + try: |
| 283 | + main() |
| 284 | + except KeyboardInterrupt: |
| 285 | + print("\n[info] Interrupted by user.") |
| 286 | + sys.exit(0) |
0 commit comments