tiny_trigger/video.py

from __future__ import annotations

import hashlib
import logging
import shutil
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Callable
from uuid import uuid4

from .detector import Detector, UltralyticsYOLOEDetector, parse_class_prompt, suppress_duplicate_detections
from .models import ActionEvent, Detection, FrameSample, VideoProcessResult


ProgressCallback = Callable[[int, int | None], None]
MP4_CODEC = "mp4v"
LOGGER = logging.getLogger(__name__)


def process_video(
    *,
    video_path: str,
    class_prompt: str | list[str],
    confidence: float = 0.25,
    frame_stride: int = 5,
    sample_interval_sec: float | None = None,
    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,
    tracking_enabled: bool = False,
    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 sample_interval_sec is not None and sample_interval_sec <= 0:
        raise ValueError("sample_interval_sec must be greater than 0.")
    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.")

    if detector is None:
        LOGGER.info(
            "Loading detector model=%s classes=%s device=%s tracking=%s",
            model_name,
            ", ".join(classes),
            device or "auto",
            tracking_enabled,
        )
        detector_started = time.perf_counter()
        detector = UltralyticsYOLOEDetector(
            class_names=classes,
            model_name=model_name,
            device=device or None,
            tracking_enabled=tracking_enabled,
        )
        LOGGER.info("Detector loaded in %.2fs", time.perf_counter() - detector_started)
    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)
    effective_frame_stride = _sampling_frame_stride(
        source_fps=source_fps,
        frame_stride=frame_stride,
        sample_interval_sec=sample_interval_sec,
    )
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
    output_fps = source_fps
    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
    latest_detections: list[Detection] = []
    LOGGER.info(
        "Processing video=%s fps=%.2f size=%sx%s sample_stride=%s max_frames=%s",
        video_path,
        source_fps,
        width,
        height,
        effective_frame_stride,
        max_frames,
    )
    try:
        while True:
            ok, frame = capture.read()
            if not ok:
                break
            frame_index += 1
            if frame_index % effective_frame_stride == 0:
                if processed_frames >= max_frames:
                    break
                timestamp_sec = frame_index / source_fps
                frames.append(FrameSample(frame_index=frame_index, timestamp_sec=timestamp_sec))
                LOGGER.info(
                    "Detecting sampled frame %s/%s source_frame=%s timestamp=%.2fs",
                    processed_frames + 1,
                    max_frames,
                    frame_index,
                    timestamp_sec,
                )
                detect_started = time.perf_counter()
                frame_detections = detector.detect(
                    frame.copy(),
                    frame_index=frame_index,
                    timestamp_sec=timestamp_sec,
                    confidence=confidence,
                    image_size=image_size,
                    max_detections=max_detections,
                )
                frame_detections = suppress_duplicate_detections(frame_detections)
                latest_detections = frame_detections
                detections.extend(latest_detections)
                processed_frames += 1
                tracked_count = sum(1 for detection in latest_detections if detection.track_id is not None)
                LOGGER.info(
                    "Detected sampled frame %s/%s in %.2fs: detections=%s tracked=%s",
                    processed_frames,
                    max_frames,
                    time.perf_counter() - detect_started,
                    len(latest_detections),
                    tracked_count,
                )
                if progress:
                    progress(processed_frames, max_frames)

            _draw_detections(frame, latest_detections)
            _write_frame(writer, _fit_frame_to_output(frame, output_size))
    finally:
        writer.release()
        capture.release()
    LOGGER.info("Finalizing annotated video %s", output_path)
    _finalize_browser_mp4(output_path)
    LOGGER.info(
        "Finished video processing: sampled_frames=%s detections=%s output=%s",
        processed_frames,
        len(detections),
        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,
        frame_stride=effective_frame_stride,
        sample_interval_sec=sample_interval_sec,
    )


def render_automation_video(
    *,
    source_video_path: str,
    detections: list[Detection],
    events: list[ActionEvent],
    frame_stride: int,
    sample_interval_sec: float | None = None,
    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 sample_interval_sec is not None and sample_interval_sec <= 0:
        raise ValueError("sample_interval_sec must be greater than 0.")
    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)
    effective_frame_stride = _sampling_frame_stride(
        source_fps=source_fps,
        frame_stride=frame_stride,
        sample_interval_sec=sample_interval_sec,
    )
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
    output_fps = source_fps
    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
    latest_detections: list[Detection] = []
    latest_events: list[ActionEvent] = []
    LOGGER.info(
        "Rendering automation video=%s fps=%.2f sample_stride=%s max_frames=%s",
        source_video_path,
        source_fps,
        effective_frame_stride,
        max_frames,
    )
    try:
        while True:
            ok, frame = capture.read()
            if not ok:
                break
            frame_index += 1
            if frame_index % effective_frame_stride == 0:
                if processed_frames >= max_frames:
                    break
                latest_detections = detections_by_frame.get(frame_index, [])
                latest_events = events_by_frame.get(frame_index, [])
                processed_frames += 1

            _draw_detections(frame, latest_detections)
            if latest_events:
                _draw_action_events(frame, latest_events)
            _write_frame(writer, _fit_frame_to_output(frame, output_size))
    finally:
        writer.release()
        capture.release()
    LOGGER.info("Finalizing automation video %s", output_path)
    _finalize_browser_mp4(output_path)
    LOGGER.info("Finished automation render: output=%s", 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 _sampling_frame_stride(
    *,
    source_fps: float,
    frame_stride: int,
    sample_interval_sec: float | None,
) -> int:
    if sample_interval_sec is None:
        return frame_stride
    return max(1, round(source_fps * sample_interval_sec))


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

    writer = cv2.VideoWriter(str(output_path), cv2.VideoWriter_fourcc(*MP4_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:v",
                "libx264",
                "-pix_fmt",
                "yuv420p",
                "-preset",
                "veryfast",
                "-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_frame(writer, frame) -> None:
    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)
        track = f" #{detection.track_id}" if detection.track_id is not None else ""
        label = f"{detection.label}{track} {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