training/VLA/frame_utils.py

"""Shared frame-sampling utilities for the VLA CoT pipeline."""
from __future__ import annotations

import base64
from io import BytesIO
from pathlib import Path
from typing import List, Tuple, Union

import cv2
import numpy as np
from PIL import Image


def sample_frames_from_mp4(
    video_path: str | Path,
    n_frames: int = 8,
    resize_short: int = 336,
    return_times: bool = False,
) -> Union[List[Image.Image], Tuple[List[Image.Image], List[float]]]:
    """Uniformly sample n_frames from an mp4, resize so short side == resize_short, return PIL RGB.

    If `return_times=True`, also returns per-frame timestamps (seconds from clip start).
    Backward compat: default behaviour returns only frames.
    """
    cap = cv2.VideoCapture(str(video_path))
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = float(cap.get(cv2.CAP_PROP_FPS)) or 30.0
    if total <= 0:
        cap.release()
        raise RuntimeError(f"bad video: {video_path}")
    idx = np.linspace(0, total - 1, n_frames).round().astype(int).tolist()
    frames: List[Image.Image] = []
    cur = 0
    wanted = set(idx)
    picked = {}
    while cap.isOpened() and len(picked) < len(idx):
        ok, frame = cap.read()
        if not ok:
            break
        if cur in wanted:
            picked[cur] = frame
        cur += 1
    cap.release()
    for i in idx:
        frame = picked.get(i, None)
        if frame is None:
            frame = next(iter(picked.values()))
        h, w = frame.shape[:2]
        scale = resize_short / min(h, w)
        nh, nw = int(round(h * scale)), int(round(w * scale))
        frame = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_AREA)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(Image.fromarray(frame))
    if return_times:
        times = [float(i) / fps for i in idx]
        return frames, times
    return frames


def uniform_frame_times(total_frames: int, n_frames: int, fps: float) -> List[float]:
    """Same index layout as sample_frames_from_mp4, but without decoding — used by
    the per-frame action label builder when we only need timestamps."""
    if total_frames <= 0 or fps <= 0:
        return [0.0] * n_frames
    idx = np.linspace(0, total_frames - 1, n_frames).round().astype(int).tolist()
    return [float(i) / float(fps) for i in idx]


def sample_frames_from_image_dir(
    image_dir: str | Path,
    n_frames: int = 8,
    resize_short: int = 336,
    fps: float = 10.0,
    return_times: bool = False,
    exts: Tuple[str, ...] = (".jpg", ".jpeg", ".png"),
) -> Union[List[Image.Image], Tuple[List[Image.Image], List[float]]]:
    """Uniformly sample n_frames from a directory of ordered image files
    (e.g. DoTA: frames/{clip}/images/000000.jpg)."""
    p = Path(image_dir)
    files = sorted([f for f in p.iterdir() if f.suffix.lower() in exts])
    if not files:
        raise RuntimeError(f"no images in {p}")
    total = len(files)
    idx = np.linspace(0, total - 1, n_frames).round().astype(int).tolist()
    frames: List[Image.Image] = []
    for i in idx:
        img = cv2.imread(str(files[i]))
        if img is None:
            img = np.zeros((resize_short, resize_short, 3), dtype=np.uint8)
        h, w = img.shape[:2]
        scale = resize_short / min(h, w)
        nh, nw = int(round(h * scale)), int(round(w * scale))
        img = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        frames.append(Image.fromarray(img))
    if return_times:
        times = [float(i) / float(fps) for i in idx]
        return frames, times
    return frames


def sample_frames(
    path: str | Path,
    n_frames: int = 8,
    resize_short: int = 336,
    return_times: bool = False,
    image_dir_fps: float = 10.0,
    frame_indices: List[int] | None = None,
) -> Union[List[Image.Image], Tuple[List[Image.Image], List[float]]]:
    """Dispatcher: mp4/mkv → video sampler; directory → image-sequence sampler.
    If `frame_indices` is provided, sample those exact frame indices (used by the
    POMDP per-frame pipeline to keep labels and frames in lockstep)."""
    p = Path(path)
    if p.is_dir():
        if frame_indices is not None:
            return sample_frames_from_image_dir_by_indices(
                p, frame_indices, resize_short=resize_short,
                fps=image_dir_fps, return_times=return_times)
        return sample_frames_from_image_dir(p, n_frames=n_frames,
                                            resize_short=resize_short,
                                            fps=image_dir_fps,
                                            return_times=return_times)
    if frame_indices is not None:
        return sample_frames_from_mp4_by_indices(
            p, frame_indices, resize_short=resize_short, return_times=return_times)
    return sample_frames_from_mp4(p, n_frames=n_frames,
                                  resize_short=resize_short,
                                  return_times=return_times)


def _resize_bgr(frame: np.ndarray, resize_short: int) -> Image.Image:
    h, w = frame.shape[:2]
    scale = resize_short / min(h, w)
    nh, nw = int(round(h * scale)), int(round(w * scale))
    frame = cv2.resize(frame, (nw, nh), interpolation=cv2.INTER_AREA)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return Image.fromarray(frame)


def sample_frames_from_mp4_by_indices(
    video_path: str | Path,
    indices: List[int],
    resize_short: int = 336,
    return_times: bool = False,
) -> Union[List[Image.Image], Tuple[List[Image.Image], List[float]]]:
    """Decode specific frame indices from an mp4."""
    cap = cv2.VideoCapture(str(video_path))
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = float(cap.get(cv2.CAP_PROP_FPS)) or 30.0
    if total <= 0:
        cap.release()
        raise RuntimeError(f"bad video: {video_path}")
    clipped = [max(0, min(total - 1, int(i))) for i in indices]
    wanted_sorted = sorted(set(clipped))
    picked: dict = {}
    cur = 0
    ptr = 0
    while cap.isOpened() and ptr < len(wanted_sorted):
        ok, frame = cap.read()
        if not ok:
            break
        while ptr < len(wanted_sorted) and cur == wanted_sorted[ptr]:
            picked[cur] = frame
            ptr += 1
        cur += 1
    cap.release()
    frames: List[Image.Image] = []
    fallback = next(iter(picked.values())) if picked else None
    for i in clipped:
        f = picked.get(i, fallback)
        frames.append(_resize_bgr(f, resize_short))
    if return_times:
        times = [float(i) / fps for i in clipped]
        return frames, times
    return frames


def sample_frames_from_image_dir_by_indices(
    image_dir: str | Path,
    indices: List[int],
    resize_short: int = 336,
    fps: float = 10.0,
    return_times: bool = False,
    exts: Tuple[str, ...] = (".jpg", ".jpeg", ".png"),
) -> Union[List[Image.Image], Tuple[List[Image.Image], List[float]]]:
    """Read specific file indices from a sorted image directory."""
    p = Path(image_dir)
    files = sorted([f for f in p.iterdir() if f.suffix.lower() in exts])
    if not files:
        raise RuntimeError(f"no images in {p}")
    total = len(files)
    clipped = [max(0, min(total - 1, int(i))) for i in indices]
    frames: List[Image.Image] = []
    for i in clipped:
        img = cv2.imread(str(files[i]))
        if img is None:
            img = np.zeros((resize_short, resize_short, 3), dtype=np.uint8)
        frames.append(_resize_bgr(img, resize_short))
    if return_times:
        times = [float(i) / float(fps) for i in clipped]
        return frames, times
    return frames


def event_anchored_indices(
    total_frames: int,
    fps: float,
    time_of_event: float | None,
    n_frames: int,
    lookback_s: float = 3.0,
    post_margin_s: float = 0.0,
) -> List[int]:
    """Compute T frame indices for POMDP-friendly per-frame labels.

    * If time_of_event is provided, sample uniformly in
      [event - lookback_s, event + post_margin_s], clipped to clip bounds.
      This puts ~`lookback_s` seconds of pre-event context in the window, which
      produces a mix of SILENT / OBSERVE / ALERT frames under the default
      (0.5s, 2.5s) thresholds.
    * Otherwise (negatives, missing event), uniform over the whole clip.
    """
    if total_frames <= 0 or fps <= 0:
        return list(range(n_frames))
    if time_of_event is None or time_of_event < 0:
        return np.linspace(0, total_frames - 1, n_frames).round().astype(int).tolist()
    end_s   = min(float(total_frames - 1) / fps, time_of_event + post_margin_s)
    start_s = max(0.0, end_s - (lookback_s + post_margin_s))
    times = np.linspace(start_s, end_s, n_frames)
    return [int(round(t * fps)) for t in times]


def pil_to_data_url(img: Image.Image, quality: int = 80) -> str:
    """Encode PIL image → data URL for the OpenAI vision API."""
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()