After story 3
This commit is contained in:
Ryan Shpeherd
@@ -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)
|
||||
+4
-4
@@ -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__":
|
||||
|
||||
+72
-14
@@ -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):
|
||||
|
||||
@@ -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())
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user