"""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()