from __future__ import annotations import hashlib import shutil import subprocess import tempfile from pathlib import Path from typing import Callable from uuid import uuid4 from .detector import Detector, UltralyticsYOLOEDetector, parse_class_prompt from .models import ActionEvent, Detection, FrameSample, VideoProcessResult ProgressCallback = Callable[[int, int | None], None] MIN_PREVIEW_FPS = 8.0 MAX_PREVIEW_FPS = 30.0 MP4_CODEC_CANDIDATES = ("avc1", "H264", "h264", "mp4v") def process_video( *, video_path: str, class_prompt: str | list[str], confidence: float = 0.25, frame_stride: int = 5, max_frames: int = 120, model_name: str = "yoloe-26s-seg.pt", image_size: int | None = None, device: str | None = None, max_detections: int | None = None, detector: Detector | None = None, output_dir: str | None = None, progress: ProgressCallback | None = None, ) -> VideoProcessResult: """Sample a video, run open-vocabulary detection, and write an annotated clip.""" try: import cv2 except ImportError as exc: # pragma: no cover - optional heavy dependency raise RuntimeError("Install opencv-python-headless to process videos.") from exc classes = parse_class_prompt(class_prompt) if not classes: raise ValueError("At least one class prompt is required.") if frame_stride < 1: raise ValueError("frame_stride must be at least 1.") if max_frames < 1: raise ValueError("max_frames must be at least 1.") if image_size is not None and image_size < 32: raise ValueError("image_size must be at least 32.") if max_detections is not None and max_detections < 1: raise ValueError("max_detections must be at least 1.") detector = detector or UltralyticsYOLOEDetector(class_names=classes, model_name=model_name, device=device or None) capture = cv2.VideoCapture(video_path) if not capture.isOpened(): raise ValueError(f"Could not open video: {video_path}") source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 30.0) width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) output_fps, repeat_count = _preview_timing(source_fps, frame_stride) output_size = _browser_frame_size(width, height) output_path = _output_path(video_path, output_dir) writer = _create_browser_mp4_writer(output_path, output_fps, output_size) if writer is None: capture.release() raise ValueError(f"Could not create annotated video: {output_path}") detections: list[Detection] = [] frames: list[FrameSample] = [] processed_frames = 0 frame_index = -1 try: while True: ok, frame = capture.read() if not ok: break frame_index += 1 if frame_index % frame_stride != 0: continue if processed_frames >= max_frames: break timestamp_sec = frame_index / source_fps frames.append(FrameSample(frame_index=frame_index, timestamp_sec=timestamp_sec)) frame_detections = detector.detect( frame, frame_index=frame_index, timestamp_sec=timestamp_sec, confidence=confidence, image_size=image_size, max_detections=max_detections, ) detections.extend(frame_detections) _draw_detections(frame, frame_detections) _write_preview_frame(writer, _fit_frame_to_output(frame, output_size), repeat_count) processed_frames += 1 if progress: progress(processed_frames, max_frames) finally: writer.release() capture.release() _finalize_browser_mp4(output_path) return VideoProcessResult( output_video_path=str(output_path), classes=classes, detections=detections, frames=frames, processed_frames=processed_frames, source_fps=source_fps, output_fps=output_fps, ) def render_automation_video( *, source_video_path: str, detections: list[Detection], events: list[ActionEvent], frame_stride: int, max_frames: int, output_dir: str | None = None, ) -> str: """Render detections plus fired automation events without rerunning inference.""" try: import cv2 except ImportError as exc: # pragma: no cover - optional heavy dependency raise RuntimeError("Install opencv-python-headless to render videos.") from exc if frame_stride < 1: raise ValueError("frame_stride must be at least 1.") if max_frames < 1: raise ValueError("max_frames must be at least 1.") capture = cv2.VideoCapture(source_video_path) if not capture.isOpened(): raise ValueError(f"Could not open video: {source_video_path}") source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 30.0) width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) output_fps, repeat_count = _preview_timing(source_fps, frame_stride) output_size = _browser_frame_size(width, height) output_path = _output_path(source_video_path, output_dir, suffix="automated") writer = _create_browser_mp4_writer(output_path, output_fps, output_size) if writer is None: capture.release() raise ValueError(f"Could not create automation video: {output_path}") detections_by_frame = _group_detections_by_frame(detections) events_by_frame = _group_events_by_frame(events) processed_frames = 0 frame_index = -1 try: while True: ok, frame = capture.read() if not ok: break frame_index += 1 if frame_index % frame_stride != 0: continue if processed_frames >= max_frames: break _draw_detections(frame, detections_by_frame.get(frame_index, [])) frame_events = events_by_frame.get(frame_index, []) if frame_events: _draw_action_events(frame, frame_events) _write_preview_frame(writer, _fit_frame_to_output(frame, output_size), repeat_count) processed_frames += 1 finally: writer.release() capture.release() _finalize_browser_mp4(output_path) return str(output_path) def _output_path(video_path: str, output_dir: str | None, *, suffix: str = "annotated") -> Path: base_dir = Path(output_dir) if output_dir else Path(tempfile.gettempdir()) / "tiny-trigger" base_dir.mkdir(parents=True, exist_ok=True) return base_dir / f"{Path(video_path).stem}-{uuid4().hex[:8]}-{suffix}.mp4" def _preview_timing(source_fps: float, frame_stride: int) -> tuple[float, int]: sampled_fps = source_fps / frame_stride if frame_stride else source_fps output_fps = min(MAX_PREVIEW_FPS, max(MIN_PREVIEW_FPS, sampled_fps or MIN_PREVIEW_FPS)) repeat_count = max(1, round(output_fps / sampled_fps)) if sampled_fps > 0 else 1 return output_fps, repeat_count def _preview_fps(source_fps: float, frame_stride: int) -> float: return _preview_timing(source_fps, frame_stride)[0] def _browser_frame_size(width: int, height: int) -> tuple[int, int]: output_width = width - (width % 2) output_height = height - (height % 2) if output_width < 2 or output_height < 2: raise ValueError("Video dimensions are too small to render.") return (output_width, output_height) def _create_browser_mp4_writer(output_path: Path, fps: float, frame_size: tuple[int, int]): import cv2 for codec in MP4_CODEC_CANDIDATES: writer = cv2.VideoWriter(str(output_path), cv2.VideoWriter_fourcc(*codec), fps, frame_size) if writer.isOpened(): return writer writer.release() return None def _finalize_browser_mp4(output_path: Path) -> None: ffmpeg_executable = _ffmpeg_executable() if not ffmpeg_executable or not output_path.exists(): return faststart_path = output_path.with_name(f"{output_path.stem}-faststart-{uuid4().hex[:8]}{output_path.suffix}") try: subprocess.run( [ ffmpeg_executable, "-y", "-loglevel", "error", "-i", str(output_path), "-c", "copy", "-movflags", "+faststart", str(faststart_path), ], check=True, capture_output=True, ) if faststart_path.exists() and faststart_path.stat().st_size > 0: faststart_path.replace(output_path) except (OSError, subprocess.CalledProcessError): if faststart_path.exists(): faststart_path.unlink() def _ffmpeg_executable() -> str | None: if ffmpeg_path := shutil.which("ffmpeg"): return ffmpeg_path try: import imageio_ffmpeg except ImportError: return None return imageio_ffmpeg.get_ffmpeg_exe() def _fit_frame_to_output(frame, output_size: tuple[int, int]): output_width, output_height = output_size height, width = frame.shape[:2] if width == output_width and height == output_height: return frame return frame[:output_height, :output_width] def _write_preview_frame(writer, frame, repeat_count: int) -> None: for _ in range(repeat_count): writer.write(frame) def _draw_detections(frame, detections: list[Detection]) -> None: import cv2 for detection in detections: x1, y1, x2, y2 = [int(value) for value in detection.bbox_xyxy] color = _color_for_label(detection.label) label = f"{detection.label} {detection.confidence:.2f}" cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) text_y = max(18, y1 - 8) cv2.putText(frame, label, (x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2, cv2.LINE_AA) def _draw_action_events(frame, events: list[ActionEvent]) -> None: import cv2 height, width = frame.shape[:2] banner_height = min(110, max(70, height // 9)) overlay = frame.copy() cv2.rectangle(overlay, (0, 0), (width, banner_height), (0, 96, 255), -1) cv2.addWeighted(overlay, 0.78, frame, 0.22, 0, frame) cv2.putText(frame, "FIRED", (24, 44), cv2.FONT_HERSHEY_SIMPLEX, 1.15, (255, 255, 255), 3, cv2.LINE_AA) details = " | ".join(f"{event.rule}: {event.action}" for event in events[:3]) cv2.putText(frame, details, (24, banner_height - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.58, (255, 255, 255), 2, cv2.LINE_AA) def _color_for_label(label: str) -> tuple[int, int, int]: digest = hashlib.md5(label.encode("utf-8")).digest() return (int(digest[0]), int(digest[1]), int(digest[2])) def _group_detections_by_frame(detections: list[Detection]) -> dict[int, list[Detection]]: grouped: dict[int, list[Detection]] = {} for detection in detections: grouped.setdefault(detection.frame_index, []).append(detection) return grouped def _group_events_by_frame(events: list[ActionEvent]) -> dict[int, list[ActionEvent]]: grouped: dict[int, list[ActionEvent]] = {} for event in events: grouped.setdefault(event.frame_index, []).append(event) return grouped