After story 3
This commit is contained in:
@@ -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 config_loader import get_config
|
||||||
from db_connector import DBConnector
|
from db_connector import DBConnector
|
||||||
from logging_config import setup_logging
|
from logging_config import setup_logging
|
||||||
|
from orchestrator import WorkerPool
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -85,14 +86,13 @@ def main():
|
|||||||
|
|
||||||
logger.info("Worker initialization complete. Starting processing loop...")
|
logger.info("Worker initialization complete. Starting processing loop...")
|
||||||
|
|
||||||
# TODO: Start scanner, processor, and monitoring services
|
pool = WorkerPool(db, config.data, max_workers=1)
|
||||||
# This is the skeleton - actual processing logic is in subsequent stories
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
while True:
|
pool.start()
|
||||||
time.sleep(60) # Main loop placeholder
|
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
logger.info("Worker shutting down.")
|
logger.info("Worker shutting down.")
|
||||||
|
pool.stop()
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
+72
-14
@@ -10,8 +10,13 @@ import time
|
|||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
from pathlib import Path
|
||||||
from typing import Dict, List, Optional
|
from typing import Dict, List, Optional
|
||||||
|
|
||||||
|
from frame_sampler import FrameSampler
|
||||||
|
from prober import VideoProber
|
||||||
|
from scratch_manager import ScratchManager
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -27,10 +32,17 @@ class JobStatus(Enum):
|
|||||||
class Job:
|
class Job:
|
||||||
"""Represents a single video processing 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.video_id = video_id
|
||||||
self.file_path = file_path
|
self.file_path = file_path
|
||||||
self.priority = priority # Higher = more urgent (based on modification time)
|
self.priority = priority # Higher = more urgent (based on modification time)
|
||||||
|
self.sampling_interval_seconds = sampling_interval_seconds
|
||||||
self.status = JobStatus.PENDING
|
self.status = JobStatus.PENDING
|
||||||
self.created_at = datetime.now(timezone.utc)
|
self.created_at = datetime.now(timezone.utc)
|
||||||
self.started_at: Optional[datetime] = None
|
self.started_at: Optional[datetime] = None
|
||||||
@@ -67,6 +79,8 @@ class WorkerPool:
|
|||||||
self._running = False
|
self._running = False
|
||||||
self._jobs_processed = 0
|
self._jobs_processed = 0
|
||||||
self._jobs_failed = 0
|
self._jobs_failed = 0
|
||||||
|
self._sampling_config = (config or {}).get("sampling", {})
|
||||||
|
self._storage_config = (config or {}).get("storage", {})
|
||||||
|
|
||||||
def start(self):
|
def start(self):
|
||||||
"""Start the worker pool."""
|
"""Start the worker pool."""
|
||||||
@@ -122,7 +136,11 @@ class WorkerPool:
|
|||||||
)
|
)
|
||||||
|
|
||||||
return [
|
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
|
for row in locked
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -135,12 +153,44 @@ class WorkerPool:
|
|||||||
logger.info("Processing job: %s (attempt %d)", job, job.attempts)
|
logger.info("Processing job: %s (attempt %d)", job, job.attempts)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# TODO: Actual processing logic (frame sampling, detection, classification)
|
sampling_interval = job.sampling_interval_seconds or self._sampling_config.get("interval_seconds", 30)
|
||||||
# This is a placeholder - will be implemented in STORY-03 through STORY-06
|
scratch_base_path = self._storage_config.get("scratch_path", "/scratch")
|
||||||
time.sleep(1) # Simulate processing
|
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
|
prober = VideoProber(timeout=10)
|
||||||
self._complete_job(job)
|
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
|
self._jobs_processed += 1
|
||||||
return True
|
return True
|
||||||
|
|
||||||
@@ -158,18 +208,26 @@ class WorkerPool:
|
|||||||
self._jobs_failed += 1
|
self._jobs_failed += 1
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def _complete_job(self, job: Job):
|
def _complete_job(self, job: Job, frame_count: Optional[int] = None):
|
||||||
"""Mark a job as completed."""
|
"""Mark a job as completed."""
|
||||||
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
now = datetime.now(timezone.utc).replace(tzinfo=None)
|
||||||
job.status = JobStatus.COMPLETED
|
job.status = JobStatus.COMPLETED
|
||||||
job.completed_at = datetime.now(timezone.utc)
|
job.completed_at = datetime.now(timezone.utc)
|
||||||
|
|
||||||
self.db.execute(
|
if frame_count is None:
|
||||||
"""UPDATE videos SET status = 'COMPLETED', updated_at = %s
|
self.db.execute(
|
||||||
WHERE id = %s""",
|
"""UPDATE videos SET status = 'COMPLETED', updated_at = %s
|
||||||
(now, job.video_id),
|
WHERE id = %s""",
|
||||||
transaction=True,
|
(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)
|
logger.info("Job completed: %s", job)
|
||||||
|
|
||||||
def _fail_job(self, job: 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