Spaces:
Sleeping
Sleeping
| 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 | |