training/Nexar/video_utils.py

#!/usr/bin/env python3
"""
Video loading utilities for Nexar mp4 clips.
Uses decord for fast frame extraction.
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
from PIL import Image

logger = logging.getLogger("Nexar.video")


def _load_with_decord(
    video_path: str,
    frame_indices: List[int],
    width: int = 640,
    height: int = 360,
) -> List[Image.Image]:
    """Extract specific frames using decord (fast)."""
    try:
        import decord
        decord.bridge.set_bridge("native")
        vr = decord.VideoReader(video_path, width=width, height=height)
        # clamp indices to valid range
        n = len(vr)
        indices = [max(0, min(idx, n - 1)) for idx in frame_indices]
        frames = vr.get_batch(indices).asnumpy()   # [N, H, W, C] uint8
        return [Image.fromarray(f) for f in frames]
    except Exception as e:
        logger.warning(f"decord failed for {video_path}: {e}; falling back to cv2")
        return _load_with_cv2(video_path, frame_indices, width, height)


def _load_with_cv2(
    video_path: str,
    frame_indices: List[int],
    width: int = 640,
    height: int = 360,
) -> List[Image.Image]:
    """Fallback: extract frames using OpenCV."""
    import cv2
    cap = cv2.VideoCapture(video_path)
    n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frames = []
    for idx in frame_indices:
        idx = max(0, min(idx, n_frames - 1))
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame)
            if width and height:
                img = img.resize((width, height), Image.LANCZOS)
            frames.append(img)
    cap.release()
    return frames


def get_video_info(video_path: str) -> Tuple[float, int]:
    """Returns (fps, n_frames)."""
    try:
        import decord
        decord.bridge.set_bridge("native")
        vr = decord.VideoReader(video_path)
        fps = vr.get_avg_fps()
        return fps, len(vr)
    except Exception:
        import cv2
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()
        return fps, n


def sample_window_frames(
    video_path: str,
    window_start_s: float,
    window_end_s: float,
    n_frames: int = 8,
    width: int = 640,
    height: int = 360,
) -> List[Image.Image]:
    """
    Extract n_frames evenly spaced from [window_start_s, window_end_s].
    Clamps to valid frame range.
    """
    fps, n_total = get_video_info(video_path)
    if fps <= 0:
        fps = 30.0

    duration = n_total / fps
    ws = max(0.0, min(window_start_s, duration))
    we = max(ws, min(window_end_s, duration))

    if we <= ws:
        we = min(ws + 0.1, duration)

    times = np.linspace(ws, we, n_frames)
    indices = [int(t * fps) for t in times]
    indices = [max(0, min(idx, n_total - 1)) for idx in indices]

    frames = _load_with_decord(video_path, indices, width, height)
    if not frames:
        frames = [Image.new("RGB", (width, height), (64, 64, 64))]
    return frames


def sample_last_window(
    video_path: str,
    window_duration_s: float = 3.0,
    n_frames: int = 8,
    width: int = 640,
    height: int = 360,
) -> List[Image.Image]:
    """
    Extract n_frames from the last `window_duration_s` seconds of the clip.
    This is the most relevant window for collision prediction (closest to event).
    """
    fps, n_total = get_video_info(video_path)
    if fps <= 0:
        fps = 30.0
    duration = n_total / fps
    window_start = max(0.0, duration - window_duration_s)
    return sample_window_frames(video_path, window_start, duration, n_frames, width, height)


def sample_multi_windows(
    video_path: str,
    n_windows: int = 3,
    window_duration_s: float = 3.0,
    n_frames_per_window: int = 8,
    width: int = 640,
    height: int = 360,
    end_offset_s: float = 0.0,
) -> List[List[Image.Image]]:
    """
    Extract n_windows temporally-spaced windows from a clip, all ending at
    `clip_end - end_offset_s`.  Windows are non-overlapping and evenly spaced.

    Returns: list of n_windows frame-lists, ordered earliest→latest.
    """
    fps, n_total = get_video_info(video_path)
    if fps <= 0:
        fps = 30.0
    duration = n_total / fps
    clip_end = duration - end_offset_s
    clip_start = max(0.0, clip_end - n_windows * window_duration_s)

    windows = []
    for i in range(n_windows):
        ws = clip_start + i * window_duration_s
        we = ws + window_duration_s
        we = min(we, clip_end)
        frames = sample_window_frames(video_path, ws, we, n_frames_per_window, width, height)
        windows.append(frames)
    return windows