From 4fdafeef916dc57cdbe92868aa8382d16479d85d Mon Sep 17 00:00:00 2001 From: Ryan Shpeherd Date: Mon, 10 Aug 2026 10:20:36 -0400 Subject: [PATCH] After story 3 --- src/frame_sampler.py | 163 +++++++++++++++++++++++++++++++++++++++++ src/main.py | 8 +- src/orchestrator.py | 86 ++++++++++++++++++---- src/scratch_manager.py | 35 +++++++++ tests/test_story_03.py | 67 +++++++++++++++++ 5 files changed, 341 insertions(+), 18 deletions(-) create mode 100644 src/frame_sampler.py create mode 100644 src/scratch_manager.py create mode 100644 tests/test_story_03.py diff --git a/src/frame_sampler.py b/src/frame_sampler.py new file mode 100644 index 0000000..8cc0029 --- /dev/null +++ b/src/frame_sampler.py @@ -0,0 +1,163 @@ +""" +Frame sampling module for VideoDetect. + +Extracts uniform temporal frames from videos using FFmpeg, handles resolution +constraints, and writes JPEG frames to a scratch directory. +""" + +import logging +import os +import subprocess +from pathlib import Path +from typing import Iterable, List, Optional, Sequence, Tuple + +logger = logging.getLogger(__name__) + + +def calculate_timestamps(duration: Optional[float], interval: float) -> List[float]: + """Return uniformly spaced timestamps in seconds for a video duration.""" + if duration is None or duration <= 0: + return [0.0] + if interval <= 0: + interval = 30.0 + + count = max(1, int(duration / interval)) + step = duration / count + return [round(i * step, 3) for i in range(count)] + + +class FrameSampler: + """Extract sampled frames from videos using FFmpeg.""" + + def __init__(self, interval_seconds: int = 30, quality: int = 2, output_format: str = "jpeg"): + self.interval_seconds = interval_seconds + self.quality = quality + self.output_format = output_format.lower() + + def extract_frame( + self, + video_path: str, + output_path: str, + timestamp: float, + resolution: Optional[Tuple[int, int]] = None, + timeout_seconds: int = 30, + ) -> Optional[str]: + """Extract a single frame at the provided timestamp.""" + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + + cmd = self._build_ffmpeg_command( + video_path=video_path, + output_path=str(output), + timestamp=timestamp, + resolution=resolution, + timeout_seconds=timeout_seconds, + ) + + try: + completed = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout_seconds + 5, + check=False, + ) + except subprocess.TimeoutExpired as exc: + logger.warning("Frame extraction timed out for %s at %.3fs", video_path, timestamp) + return None + + if completed.returncode != 0: + logger.warning( + "Frame extraction failed for %s at %.3fs: %s", + video_path, + timestamp, + completed.stderr.strip(), + ) + return None + + if output.exists(): + return str(output) + return None + + def extract_frames( + self, + video_path: str, + output_dir: str, + duration: Optional[float] = None, + interval_seconds: Optional[int] = None, + resolution: Optional[Tuple[int, int]] = None, + timeout_seconds: int = 30, + ) -> List[str]: + """Extract uniform frames for the video and return output paths.""" + resolved_interval = interval_seconds or self.interval_seconds + timestamps = calculate_timestamps(duration=duration, interval=float(resolved_interval)) + + extracted: List[str] = [] + for ts in timestamps: + frame_name = self._frame_name_from_timestamp(ts, video_path) + output_path = str(Path(output_dir) / frame_name) + result = self.extract_frame( + video_path=video_path, + output_path=output_path, + timestamp=ts, + resolution=resolution, + timeout_seconds=timeout_seconds, + ) + if result: + extracted.append(result) + + return extracted + + def _build_ffmpeg_command( + self, + video_path: str, + output_path: str, + timestamp: float, + resolution: Optional[Tuple[int, int]] = None, + timeout_seconds: int = 30, + ) -> List[str]: + """Build an FFmpeg command for extracting a single frame.""" + cmd = ["ffmpeg", "-y", "-ss", str(timestamp), "-i", video_path] + + vf_filters = [] + if resolution is not None: + target_width, target_height = self._resolve_target_resolution(resolution) + vf_filters.append(f"scale={target_width}:{target_height}:force_original_aspect_ratio=decrease") + + if vf_filters: + cmd.extend(["-vf", ",".join(vf_filters)]) + + cmd.extend([ + "-vframes", + "1", + "-q:v", + str(self.quality), + "-f", + self.output_format, + output_path, + ]) + return cmd + + @staticmethod + def _resolve_target_resolution(resolution: Tuple[int, int]) -> Tuple[int, int]: + """Downscale frames larger than 1080p to fit VRAM constraints.""" + width, height = resolution + if max(width, height) > 1080: + if width >= 3840 or height >= 2160: + return 1920, 1080 + if width >= 2560 or height >= 1440: + return 1280, 720 + return width, height + + @staticmethod + def _frame_name_from_timestamp(timestamp: float, video_path: str) -> str: + video_id = Path(video_path).stem + timestamp_ms = int(round(timestamp * 1000)) + return f"{video_id}_{timestamp_ms}.jpg" + + +def get_target_resolution(resolution: Optional[Tuple[int, int]]) -> Tuple[int, int]: + """Compatibility helper for resolution handling.""" + if resolution is None: + return (0, 0) + return FrameSampler._resolve_target_resolution(resolution) diff --git a/src/main.py b/src/main.py index 0480276..ef7ee9f 100644 --- a/src/main.py +++ b/src/main.py @@ -17,6 +17,7 @@ sys.path.insert(0, str(Path(__file__).parent)) from config_loader import get_config from db_connector import DBConnector from logging_config import setup_logging +from orchestrator import WorkerPool logger = logging.getLogger(__name__) @@ -85,14 +86,13 @@ def main(): logger.info("Worker initialization complete. Starting processing loop...") - # TODO: Start scanner, processor, and monitoring services - # This is the skeleton - actual processing logic is in subsequent stories + pool = WorkerPool(db, config.data, max_workers=1) try: - while True: - time.sleep(60) # Main loop placeholder + pool.start() except KeyboardInterrupt: logger.info("Worker shutting down.") + pool.stop() if __name__ == "__main__": diff --git a/src/orchestrator.py b/src/orchestrator.py index 9b6084a..2450bf9 100644 --- a/src/orchestrator.py +++ b/src/orchestrator.py @@ -10,8 +10,13 @@ import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from enum import Enum +from pathlib import Path from typing import Dict, List, Optional +from frame_sampler import FrameSampler +from prober import VideoProber +from scratch_manager import ScratchManager + logger = logging.getLogger(__name__) @@ -27,10 +32,17 @@ class JobStatus(Enum): class Job: """Represents a single video processing job.""" - def __init__(self, video_id: int, file_path: str, priority: float = 0.0): + def __init__( + self, + video_id: int, + file_path: str, + priority: float = 0.0, + sampling_interval_seconds: Optional[int] = None, + ): self.video_id = video_id self.file_path = file_path self.priority = priority # Higher = more urgent (based on modification time) + self.sampling_interval_seconds = sampling_interval_seconds self.status = JobStatus.PENDING self.created_at = datetime.now(timezone.utc) self.started_at: Optional[datetime] = None @@ -67,6 +79,8 @@ class WorkerPool: self._running = False self._jobs_processed = 0 self._jobs_failed = 0 + self._sampling_config = (config or {}).get("sampling", {}) + self._storage_config = (config or {}).get("storage", {}) def start(self): """Start the worker pool.""" @@ -122,7 +136,11 @@ class WorkerPool: ) return [ - Job(video_id=row["id"], file_path=row["file_path"], priority=row["last_scan_time"].timestamp()) + Job( + video_id=row["id"], + file_path=row["file_path"], + priority=row["last_scan_time"].timestamp(), + ) for row in locked ] @@ -135,12 +153,44 @@ class WorkerPool: logger.info("Processing job: %s (attempt %d)", job, job.attempts) try: - # TODO: Actual processing logic (frame sampling, detection, classification) - # This is a placeholder - will be implemented in STORY-03 through STORY-06 - time.sleep(1) # Simulate processing + sampling_interval = job.sampling_interval_seconds or self._sampling_config.get("interval_seconds", 30) + scratch_base_path = self._storage_config.get("scratch_path", "/scratch") + scratch_manager = ScratchManager( + base_path=scratch_base_path, + video_id=str(job.video_id), + auto_cleanup=True, + ) + frame_dir = scratch_manager.ensure_frame_dir() - # Mark as completed - self._complete_job(job) + prober = VideoProber(timeout=10) + metadata = prober.probe(job.file_path) + if metadata.is_unscannable or metadata.duration is None: + raise RuntimeError(metadata.error_message or "Video metadata could not be determined") + + resolution = None + if metadata.resolution_w and metadata.resolution_h: + resolution = (metadata.resolution_w, metadata.resolution_h) + + sampler = FrameSampler( + interval_seconds=int(sampling_interval), + quality=int(self._sampling_config.get("quality", 2)), + output_format=self._sampling_config.get("format", "jpeg"), + ) + extracted_frames = sampler.extract_frames( + video_path=job.file_path, + output_dir=str(frame_dir), + duration=metadata.duration, + interval_seconds=int(sampling_interval), + resolution=resolution, + timeout_seconds=30, + ) + + if not extracted_frames: + scratch_manager.cleanup() + raise RuntimeError("No frames were extracted") + + self._complete_job(job, frame_count=len(extracted_frames)) + scratch_manager.cleanup() self._jobs_processed += 1 return True @@ -158,18 +208,26 @@ class WorkerPool: self._jobs_failed += 1 return False - def _complete_job(self, job: Job): + def _complete_job(self, job: Job, frame_count: Optional[int] = None): """Mark a job as completed.""" now = datetime.now(timezone.utc).replace(tzinfo=None) job.status = JobStatus.COMPLETED job.completed_at = datetime.now(timezone.utc) - self.db.execute( - """UPDATE videos SET status = 'COMPLETED', updated_at = %s - WHERE id = %s""", - (now, job.video_id), - transaction=True, - ) + if frame_count is None: + self.db.execute( + """UPDATE videos SET status = 'COMPLETED', updated_at = %s + WHERE id = %s""", + (now, job.video_id), + transaction=True, + ) + else: + self.db.execute( + """UPDATE videos SET status = 'COMPLETED', frame_count = %s, updated_at = %s + WHERE id = %s""", + (frame_count, now, job.video_id), + transaction=True, + ) logger.info("Job completed: %s", job) def _fail_job(self, job: Job): diff --git a/src/scratch_manager.py b/src/scratch_manager.py new file mode 100644 index 0000000..7d6fc1f --- /dev/null +++ b/src/scratch_manager.py @@ -0,0 +1,35 @@ +"""Utilities for managing scratch space used by frame sampling.""" + +import logging +import shutil +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + + +class ScratchManager: + """Create and clean per-video frame scratch directories.""" + + def __init__(self, base_path: str = "/scratch", video_id: Optional[str] = None, auto_cleanup: bool = True): + self.base_path = Path(base_path) + self.video_id = video_id or "unknown" + self.auto_cleanup = auto_cleanup + self.frame_dir = self.base_path / self.video_id / "frames" + + def ensure_frame_dir(self) -> Path: + """Create the per-video scratch directory if it does not exist.""" + self.frame_dir.mkdir(parents=True, exist_ok=True) + return self.frame_dir + + def cleanup(self) -> None: + """Remove the video-specific frame directory and everything underneath it.""" + if self.auto_cleanup and self.frame_dir.exists(): + shutil.rmtree(self.frame_dir) + logger.info("Cleaned scratch frames for %s", self.video_id) + + def usage_bytes(self) -> int: + """Return the total size of the scratch directory in bytes.""" + if not self.frame_dir.exists(): + return 0 + return sum(path.stat().st_size for path in self.frame_dir.rglob("*") if path.is_file()) diff --git a/tests/test_story_03.py b/tests/test_story_03.py new file mode 100644 index 0000000..c8b15a3 --- /dev/null +++ b/tests/test_story_03.py @@ -0,0 +1,67 @@ +import shutil +import sys +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) + +from frame_sampler import FrameSampler, calculate_timestamps +from scratch_manager import ScratchManager + + +class Story03SamplingTests(unittest.TestCase): + def test_calculate_timestamps_uses_uniform_temporal_spacing(self): + stamps = calculate_timestamps(duration=90.0, interval=30.0) + self.assertEqual(stamps, [0.0, 30.0, 60.0]) + + def test_frame_sampler_builds_ffmpeg_command_with_scale_and_jpeg_output(self): + sampler = FrameSampler(interval_seconds=30, quality=2, output_format="jpeg") + command = sampler._build_ffmpeg_command( + video_path="/tmp/video.mp4", + output_path="/tmp/out.jpg", + timestamp=12.5, + resolution=(3840, 2160), + ) + + self.assertIn("ffmpeg", command[0]) + self.assertIn("-ss", command) + self.assertIn("-vframes", command) + self.assertIn("scale=1920:1080", " ".join(command)) + self.assertTrue(command[-1].endswith("out.jpg")) + + def test_frame_sampler_uses_subprocess_and_returns_output_path(self): + sampler = FrameSampler(interval_seconds=30, quality=2, output_format="jpeg") + with tempfile.TemporaryDirectory() as tmpdir: + video_path = Path(tmpdir) / "sample.mp4" + output_path = Path(tmpdir) / "frame.jpg" + video_path.write_bytes(b"fake") + + def _mock_run(*args, **kwargs): + output_path.write_bytes(b"frame") + return type("Completed", (), {"returncode": 0, "stdout": b"", "stderr": b""})() + + with patch("subprocess.run", side_effect=_mock_run): + result = sampler.extract_frame( + video_path=str(video_path), + output_path=str(output_path), + timestamp=15.0, + resolution=(1280, 720), + ) + + self.assertTrue(result) + self.assertEqual(output_path.name, Path(result).name) + + def test_scratch_manager_cleans_up_frames_after_processing(self): + with tempfile.TemporaryDirectory() as tmpdir: + manager = ScratchManager(base_path=tmpdir, video_id="video-1", auto_cleanup=True) + frame_dir = manager.ensure_frame_dir() + (frame_dir / "video-1_1000.jpg").write_bytes(b"frame") + + manager.cleanup() + self.assertFalse(frame_dir.exists()) + + +if __name__ == "__main__": + unittest.main()