Blame for tiny_trigger/video.py

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