preprocessing.py

"""
preprocessing.py
================
Faithful port of the *always-on* training-time video preprocessing used by
`mamounyosef/sign-language-bridge`.

The adapter was trained with three preprocessing stages applied to **every**
split (train / val / test), so inference must reproduce them or the model sees
an out-of-distribution input distribution and quality collapses:

  1. Pose-guided signer crop  — MediaPipe PoseLandmarker, union of upper-body +
     hand landmarks across the clip, 25 % padding, snapped down to a multiple
     of 32 (Qwen3-VL patch_size 16 x merge_size 2).
  2. CLAHE                    — on the L channel in LAB, clipLimit 2.0, 8x8 tiles.
  3. Landmark overlay         — RTMPose Wholebody (COCO-Wholebody 133 kpts):
     6 upper-body joints + 21 keypoints per hand, drawn as a 1 px skeleton.

Sources this mirrors, in the project repo
(https://github.com/mamounyosef/sign-language-bridge):

  * `data_code/signer_cropper.py`                      -> SignerCropper
  * `data_code/21_extract_landmarks.py`                -> LandmarkExtractor, _postprocess
  * `model_training_scripts/qwen3vl_training.py`
        Qwen3VLCollator._apply_signer_crop             -> apply_signer_crop
        Qwen3VLCollator._apply_clahe_opencv            -> apply_clahe
        Qwen3VLCollator._apply_landmark_overlay        -> apply_landmark_overlay

One deliberate deviation: the training pipeline extracted landmarks offline
from *native-resolution* crops and mapped them onto the decoded frames by
nearest native frame index. Here the landmarks are extracted directly from the
decoded (already pixel-budgeted) crop, so the mapping is 1:1. Landmark
coordinates are normalized to the crop in both cases, so the drawn overlay is
equivalent; this just removes a resampling step.
"""

from __future__ import annotations

import logging
import os

# Silence MediaPipe / glog native-layer INFO + WARNING spam BEFORE importing
# mediapipe -- these env vars are read when the native library loads.
os.environ.setdefault("GLOG_minloglevel", "2")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
os.environ.setdefault("absl_logging_verbosity", "2")

from dataclasses import dataclass
from typing import Iterable, Optional

import cv2
import numpy as np
import torch

logger = logging.getLogger(__name__)

# Qwen3-VL requires H/W divisible by patch_size * merge_size = 16 * 2 = 32.
SNAP_ALIGN = 32


# ============================================================================
#  1. Pose-guided signer crop  (port of data_code/signer_cropper.py)
# ============================================================================

# Upper-body + hand landmarks (excludes legs/hips for a tight signing-space box):
#   0 nose | 1-6 eyes | 7-8 ears | 9-10 mouth | 11-12 shoulders
#   13-14 elbows | 15-16 wrists | 17-22 hand keypoints
SIGNING_LANDMARK_INDICES: tuple[int, ...] = tuple(range(0, 23))
MIN_LANDMARK_VISIBILITY = 0.5

POSE_MODEL_URLS = {
    "lite": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_lite/float16/latest/pose_landmarker_lite.task",
    "full": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_full/float16/latest/pose_landmarker_full.task",
    "heavy": "https://storage.googleapis.com/mediapipe-models/pose_landmarker/pose_landmarker_heavy/float16/latest/pose_landmarker_heavy.task",
}


@dataclass(frozen=True)
class BBoxResult:
    """Signer bbox in PIXEL space of the source video (x = width, y = height)."""

    x1: int
    y1: int
    x2: int
    y2: int
    frame_width: int
    frame_height: int
    detection_rate: float
    num_sampled_frames: int
    num_detected_frames: int
    failed: bool


class SignerCropper:
    """One static bbox per clip: union of pose landmarks over sampled frames.

    Using the union (rather than a per-frame box) guarantees a briefly extended
    hand stays inside the crop for the whole clip, which keeps the signing space
    stable -- exactly what the model was trained on.
    """

    def __init__(
        self,
        models_dir: str,
        sample_every_n: int = 4,
        padding_ratio: float = 0.25,
        min_detection_confidence: float = 0.5,
        min_presence_confidence: float = 0.5,
        landmark_indices: Iterable[int] = SIGNING_LANDMARK_INDICES,
        min_landmark_visibility: float = MIN_LANDMARK_VISIBILITY,
        model_variant: str = "full",
    ) -> None:
        self.sample_every_n = int(sample_every_n)
        self.padding_ratio = float(padding_ratio)
        self.landmark_indices = tuple(landmark_indices)
        self.min_landmark_visibility = float(min_landmark_visibility)

        import mediapipe as mp
        from mediapipe.tasks import python as mp_python
        from mediapipe.tasks.python import vision as mp_vision

        self._mp = mp
        self._mp_vision = mp_vision

        model_path = _download_once(
            POSE_MODEL_URLS[model_variant],
            os.path.join(models_dir, f"pose_landmarker_{model_variant}.task"),
        )
        # IMAGE mode is the default; setting running_mode / min_tracking_confidence
        # explicitly has caused proto-binding mismatches in some MediaPipe builds.
        self._options = mp_vision.PoseLandmarkerOptions(
            base_options=mp_python.BaseOptions(model_asset_path=model_path),
            num_poses=1,
            min_pose_detection_confidence=float(min_detection_confidence),
            min_pose_presence_confidence=float(min_presence_confidence),
        )

    def compute_bbox(self, frames_bgr: np.ndarray) -> BBoxResult:
        """frames_bgr: (T, H, W, 3) uint8 BGR (cv2 order). MediaPipe wants RGB."""
        T, H, W, _ = frames_bgr.shape

        sampled = list(range(0, T, self.sample_every_n))
        if sampled and sampled[-1] != T - 1:
            sampled.append(T - 1)
        elif not sampled:
            sampled = [0]

        xs: list[float] = []
        ys: list[float] = []
        num_detected = 0

        with self._mp_vision.PoseLandmarker.create_from_options(self._options) as landmarker:
            for idx in sampled:
                frame_rgb = cv2.cvtColor(frames_bgr[idx], cv2.COLOR_BGR2RGB)
                mp_image = self._mp.Image(
                    image_format=self._mp.ImageFormat.SRGB,
                    data=np.ascontiguousarray(frame_rgb),
                )
                result = landmarker.detect(mp_image)
                if not result.pose_landmarks:
                    continue
                landmarks = result.pose_landmarks[0]  # num_poses=1

                got_any = False
                for lm_idx in self.landmark_indices:
                    lm = landmarks[lm_idx]
                    visibility = getattr(lm, "visibility", 1.0)
                    if visibility is not None and visibility < self.min_landmark_visibility:
                        continue
                    # MediaPipe can predict outside 0..1 when the body is occluded.
                    if not (0.0 <= lm.x <= 1.0 and 0.0 <= lm.y <= 1.0):
                        continue
                    xs.append(lm.x * W)
                    ys.append(lm.y * H)
                    got_any = True
                if got_any:
                    num_detected += 1

        num_sampled = len(sampled)
        detection_rate = num_detected / num_sampled if num_sampled else 0.0

        def _failed() -> BBoxResult:
            return BBoxResult(0, 0, 0, 0, W, H, detection_rate, num_sampled, num_detected, True)

        if not xs:
            return _failed()

        raw_x1, raw_x2 = min(xs), max(xs)
        raw_y1, raw_y2 = min(ys), max(ys)

        pad_x = self.padding_ratio * (raw_x2 - raw_x1)
        pad_y = self.padding_ratio * (raw_y2 - raw_y1)

        x1 = int(round(max(0.0, raw_x1 - pad_x)))
        y1 = int(round(max(0.0, raw_y1 - pad_y)))
        x2 = int(round(min(float(W), raw_x2 + pad_x)))
        y2 = int(round(min(float(H), raw_y2 + pad_y)))

        if x2 <= x1 or y2 <= y1:
            return _failed()

        return BBoxResult(x1, y1, x2, y2, W, H, detection_rate, num_sampled, num_detected, False)


def _snap_down_to_align(lo: int, hi: int, align: int) -> tuple[int, int]:
    """Shrink [lo, hi) so its length is a multiple of `align`, trimming symmetrically."""
    length = hi - lo
    target = (length // align) * align
    if target <= 0 or target == length:
        return lo, hi
    trim = length - target
    left = trim // 2
    return lo + left, lo + left + target


def apply_signer_crop(video: torch.Tensor, bbox: Optional[BBoxResult]) -> torch.Tensor:
    """Crop a decoded (T, C, H, W) clip to its signer bbox.

    The bbox is in source-video pixels; `process_vision_info` returns frames at
    a different resolution, so it is rescaled to the current H/W before slicing
    (mirrors Qwen3VLCollator._apply_signer_crop).
    """
    cur_h, cur_w = int(video.shape[-2]), int(video.shape[-1])

    cx1, cy1, cx2, cy2 = 0, 0, cur_w, cur_h
    bbox_used = False

    if bbox is not None and not bbox.failed:
        sx = cur_w / max(bbox.frame_width, 1)
        sy = cur_h / max(bbox.frame_height, 1)
        bx1 = max(0, int(round(bbox.x1 * sx)))
        by1 = max(0, int(round(bbox.y1 * sy)))
        bx2 = min(cur_w, int(round(bbox.x2 * sx)))
        by2 = min(cur_h, int(round(bbox.y2 * sy)))
        if bx2 - bx1 >= SNAP_ALIGN and by2 - by1 >= SNAP_ALIGN:
            cx1, cy1, cx2, cy2 = bx1, by1, bx2, by2
            bbox_used = True
        else:
            logger.warning("Degenerate signer bbox -> falling back to the full frame.")

    cx1, cx2 = _snap_down_to_align(cx1, cx2, SNAP_ALIGN)
    cy1, cy2 = _snap_down_to_align(cy1, cy2, SNAP_ALIGN)

    take_full_frame = not bbox_used and cx1 == 0 and cy1 == 0 and cx2 == cur_w and cy2 == cur_h
    if take_full_frame:
        return video
    return video[..., cy1:cy2, cx1:cx2].contiguous()


# ============================================================================
#  2. CLAHE  (port of Qwen3VLCollator._apply_clahe_opencv)
# ============================================================================

def apply_clahe(video: torch.Tensor, clip_limit: float = 2.0,
                tile_grid: tuple[int, int] = (8, 8)) -> torch.Tensor:
    """CLAHE on the L channel in LAB. Input/output: (T, C, H, W) uint8 RGB."""
    clahe = cv2.createCLAHE(clipLimit=float(clip_limit), tileGridSize=tuple(tile_grid))
    out = video.clone()
    for t in range(video.shape[0]):
        frame_rgb = video[t].permute(1, 2, 0).contiguous().numpy()
        lab = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2LAB)
        lab[:, :, 0] = clahe.apply(lab[:, :, 0])
        frame_rgb = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
        out[t] = torch.from_numpy(frame_rgb).permute(2, 0, 1)
    return out


# ============================================================================
#  3. Landmark overlay  (port of data_code/21_extract_landmarks.py + collator)
# ============================================================================

# COCO-Wholebody 133-keypoint layout.
POSE_INDICES = [5, 6, 7, 8, 9, 10]      # L/R shoulder, L/R elbow, L/R wrist
LEFT_HAND_SLICE = slice(91, 112)        # 21 keypoints
RIGHT_HAND_SLICE = slice(112, 133)      # 21 keypoints
POSE_CONF_THRESHOLD = 0.3               # joints below this -> NaN
HAND_CONF_THRESHOLD = 0.2               # whole hand rejected if mean conf below this

# Pose array stores COCO indices [5,6,7,8,9,10] at positions [0..5].
POSE_CONNECTIONS = [
    (0, 1),  # left shoulder - right shoulder
    (0, 2),  # left shoulder - left elbow
    (2, 4),  # left elbow    - left wrist
    (1, 3),  # right shoulder- right elbow
    (3, 5),  # right elbow   - right wrist
]

HAND_CONNECTIONS = [
    (0, 1), (1, 2), (2, 3), (3, 4),            # thumb
    (0, 5), (5, 6), (6, 7), (7, 8),            # index
    (0, 9), (9, 10), (10, 11), (11, 12),       # middle
    (0, 13), (13, 14), (14, 15), (15, 16),     # ring
    (0, 17), (17, 18), (18, 19), (19, 20),     # pinky
    (5, 9), (9, 13), (13, 17),                 # palm arch
]

# BGR, as drawn by the training collator.
COLOR_POSE = (0, 255, 255)        # yellow
COLOR_LEFT_HAND = (0, 200, 0)     # green
COLOR_RIGHT_HAND = (255, 80, 0)   # blue


def _add_torch_cuda_dlls() -> None:
    """Point the Windows DLL loader at PyTorch's bundled cuDNN/CUDA libraries.

    onnxruntime-gpu needs cudnn64_9.dll, which isn't on PATH by default. PyTorch
    ships cuDNN 9, so adding its lib directory before the first ORT session is
    created is enough. No-op off Windows or without torch. (Ported from
    `data_code/21_extract_landmarks.py` in the project repo, which hit exactly
    this on the training machine.)
    """
    if os.name != "nt":
        return
    try:
        import torch as _torch

        torch_lib = os.path.join(os.path.dirname(_torch.__file__), "lib")
        if os.path.isdir(torch_lib):
            os.add_dll_directory(torch_lib)
    except Exception:  # noqa: BLE001
        pass


def _onnx_providers_available() -> list[str]:
    try:
        import onnxruntime as ort

        return list(ort.get_available_providers())
    except Exception:  # noqa: BLE001
        return []


class LandmarkExtractor:
    """RTMPose Wholebody keypoints, normalized to the crop.

    The training pipeline skips person detection: frames are already tightly
    cropped around the signer, so `bboxes=[]` makes RTMPose treat the full frame
    as the person region -- identical to what a detector would return, minus an
    entire forward pass per frame.

    The `performance` model is an x-large 384x288 network: roughly 1 s per frame
    on CPU, which makes a short clip take minutes. `device="auto"` therefore
    prefers CUDA and only falls back to CPU when no GPU execution provider is
    present.
    """

    def __init__(self, mode: str = "performance", device: str = "auto",
                 backend: str = "onnxruntime") -> None:
        from rtmlib import RTMPose, Wholebody

        if device == "auto":
            _add_torch_cuda_dlls()
            providers = _onnx_providers_available()
            device = "cuda" if "CUDAExecutionProvider" in providers else "cpu"
            if device == "cpu":
                logger.warning(
                    "No CUDA execution provider for onnxruntime (%s); landmark "
                    "extraction will run on CPU and be slow.", providers,
                )
        elif device == "cuda":
            _add_torch_cuda_dlls()

        self.device = device
        mode_cfg = Wholebody.MODE[mode]
        self.model = RTMPose(
            mode_cfg["pose"],
            model_input_size=mode_cfg["pose_input_size"],
            to_openpose=False,
            backend=backend,
            device=device,
        )

    def extract(self, frames_bgr: np.ndarray, crop_w: int, crop_h: int):
        """frames_bgr: (T, H, W, 3) uint8 BGR. Returns pose/lh/rh, NaN where absent."""
        nan6 = np.full((6, 2), np.nan, dtype=np.float32)
        nan21 = np.full((21, 2), np.nan, dtype=np.float32)

        pose_rows, lh_rows, rh_rows = [], [], []
        for frame_bgr in frames_bgr:
            try:
                keypoints, scores = self.model(np.ascontiguousarray(frame_bgr), bboxes=[])
            except Exception as exc:  # noqa: BLE001 - a bad frame must not kill the clip
                logger.warning("RTMPose failed on a frame: %r", exc)
                keypoints = []

            if len(keypoints) == 0:
                pose_rows.append(nan6.copy())
                lh_rows.append(nan21.copy())
                rh_rows.append(nan21.copy())
                continue

            kps = keypoints[0].astype(np.float32)   # (133, 2) pixel coords
            scr = scores[0].astype(np.float32)      # (133,)

            kps_norm = kps.copy()
            kps_norm[:, 0] = np.clip(kps_norm[:, 0] / max(crop_w, 1), 0.0, 1.0)
            kps_norm[:, 1] = np.clip(kps_norm[:, 1] / max(crop_h, 1), 0.0, 1.0)

            pose_row = nan6.copy()
            for out_i, coco_i in enumerate(POSE_INDICES):
                if scr[coco_i] >= POSE_CONF_THRESHOLD:
                    pose_row[out_i] = kps_norm[coco_i]
            pose_rows.append(pose_row)

            # Accept a hand on mean confidence, then NaN out weak individual joints.
            for sl, sink in ((LEFT_HAND_SLICE, lh_rows), (RIGHT_HAND_SLICE, rh_rows)):
                hand_scores = scr[sl]
                if np.mean(hand_scores) >= HAND_CONF_THRESHOLD:
                    row = kps_norm[sl].copy()
                    row[hand_scores < HAND_CONF_THRESHOLD] = np.nan
                else:
                    row = nan21.copy()
                sink.append(row)

        return (
            np.stack(pose_rows).astype(np.float32),
            np.stack(lh_rows).astype(np.float32),
            np.stack(rh_rows).astype(np.float32),
        )


def _hand_centroid(hand: np.ndarray) -> Optional[np.ndarray]:
    valid = hand[~np.any(np.isnan(hand), axis=1)]
    return valid.mean(axis=0) if len(valid) else None


def _reject_spikes(arr: np.ndarray, threshold: float = 0.12) -> np.ndarray:
    """NaN out frames whose centroid deviates from its neighbours' midpoint.

    Catches brief 'teleport' detections. threshold = 12 % of the crop dimension.
    """
    out = arr.copy()
    for t in range(1, arr.shape[0] - 1):
        c_prev = _hand_centroid(arr[t - 1])
        c_curr = _hand_centroid(arr[t])
        c_next = _hand_centroid(arr[t + 1])
        if c_prev is None or c_curr is None or c_next is None:
            continue
        if np.linalg.norm(c_curr - (c_prev + c_next) / 2.0) > threshold:
            out[t] = np.nan
    return out


def _fill_landmark_gaps(arr: np.ndarray, max_gap: int = 8) -> np.ndarray:
    """Forward-fill NaN runs of <= max_gap frames per joint (brief occlusions)."""
    arr = arr.copy()
    _, K, _ = arr.shape
    for k in range(K):
        last_valid = None
        gap_start = None
        for t in range(arr.shape[0]):
            if not np.any(np.isnan(arr[t, k])):
                if gap_start is not None and last_valid is not None and (t - gap_start) <= max_gap:
                    arr[gap_start:t, k] = last_valid
                gap_start = None
                last_valid = arr[t, k].copy()
            elif gap_start is None:
                gap_start = t
    return arr


def _smooth_landmarks(arr: np.ndarray) -> np.ndarray:
    """3-frame weighted smoothing [0.25, 0.5, 0.25], only on fully-valid windows."""
    if arr.shape[0] < 3:
        return arr
    out = arr.copy()
    for t in range(1, arr.shape[0] - 1):
        window = arr[t - 1: t + 2]
        if not np.any(np.isnan(window)):
            out[t] = 0.25 * arr[t - 1] + 0.5 * arr[t] + 0.25 * arr[t + 1]
    return out


def postprocess_landmarks(pose: np.ndarray, lh: np.ndarray, rh: np.ndarray):
    """Spike rejection -> gap fill -> temporal smoothing."""
    lh = _smooth_landmarks(_fill_landmark_gaps(_reject_spikes(lh), max_gap=8))
    rh = _smooth_landmarks(_fill_landmark_gaps(_reject_spikes(rh), max_gap=8))
    pose = _smooth_landmarks(_fill_landmark_gaps(pose, max_gap=8))
    return pose, lh, rh


def apply_landmark_overlay(video: torch.Tensor, pose: np.ndarray,
                           lh: np.ndarray, rh: np.ndarray) -> torch.Tensor:
    """Draw the pose + hand skeleton onto (T, C, H, W) uint8 RGB frames.

    Landmarks are normalized to the crop, so they scale to any resolution:
    px = norm_x * crop_W. Line/dot size is 1 px, matching training.
    """
    T_frames, _, cur_h, cur_w = video.shape
    out = video.clone()

    for t in range(T_frames):
        idx = min(t, pose.shape[0] - 1)

        frame_rgb = video[t].permute(1, 2, 0).contiguous().numpy().copy()
        frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)

        def to_px(nx, ny):
            return (int(round(float(nx) * cur_w)), int(round(float(ny) * cur_h)))

        pose_t = pose[idx]
        for i, j in POSE_CONNECTIONS:
            if np.any(np.isnan(pose_t[i])) or np.any(np.isnan(pose_t[j])):
                continue
            cv2.line(frame_bgr, to_px(*pose_t[i]), to_px(*pose_t[j]),
                     COLOR_POSE, 1, cv2.LINE_AA)
        for i in range(6):
            if not np.any(np.isnan(pose_t[i])):
                cv2.circle(frame_bgr, to_px(*pose_t[i]), 1, COLOR_POSE, -1, cv2.LINE_AA)

        for hand_arr, color in ((lh, COLOR_LEFT_HAND), (rh, COLOR_RIGHT_HAND)):
            hand = hand_arr[idx]
            if np.all(np.isnan(hand)):
                continue
            for i, j in HAND_CONNECTIONS:
                if np.any(np.isnan(hand[i])) or np.any(np.isnan(hand[j])):
                    continue
                cv2.line(frame_bgr, to_px(*hand[i]), to_px(*hand[j]), color, 1, cv2.LINE_AA)
            for i in range(21):
                if not np.any(np.isnan(hand[i])):
                    cv2.circle(frame_bgr, to_px(*hand[i]), 1, color, -1, cv2.LINE_AA)

        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        out[t] = torch.from_numpy(frame_rgb).permute(2, 0, 1)

    return out


# ============================================================================
#  Helpers
# ============================================================================

def _download_once(url: str, dest: str) -> str:
    """Download `url` to `dest` unless it is already there. Returns `dest`."""
    import urllib.request

    os.makedirs(os.path.dirname(dest), exist_ok=True)
    if not os.path.exists(dest):
        logger.info("Downloading %s -> %s", url, dest)
        tmp = dest + ".part"
        urllib.request.urlretrieve(url, tmp)
        os.replace(tmp, dest)
    return dest


def resize_video(video: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
    """Downscale a (T, C, H, W) uint8 clip, one frame at a time.

    Deliberately not `F.interpolate`: that needs a float32 copy of the whole clip
    (~200 MB for a few seconds of video) before it produces anything, which is
    enough to fail on a memory-constrained machine. cv2 works per frame in uint8,
    and INTER_AREA is the correct filter for downscaling.
    """
    T, C = video.shape[0], video.shape[1]
    out = torch.empty((T, C, target_h, target_w), dtype=torch.uint8)
    for t in range(T):
        frame = video[t].permute(1, 2, 0).contiguous().numpy()
        resized = cv2.resize(frame, (target_w, target_h), interpolation=cv2.INTER_AREA)
        out[t] = torch.from_numpy(resized).permute(2, 0, 1)
    return out


def probe_duration_seconds(path: str) -> float:
    """Clip duration in seconds from container metadata, or 0.0 if unknown."""
    cap = cv2.VideoCapture(path)
    try:
        if not cap.isOpened():
            return 0.0
        fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
        total = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
        return float(total / fps) if fps > 0 and total > 0 else 0.0
    finally:
        cap.release()


def read_video_frames_bgr(
    path: str,
    stride: int = 1,
    max_seconds: Optional[float] = None,
    max_frames: int = 2000,
) -> tuple[np.ndarray, float, int]:
    """Load every `stride`-th frame as (T, H, W, 3) uint8 BGR.

    Only the kept frames are held in memory: decoding a 1080p clip in full would
    cost ~2 GB, and the bbox pass only ever looks at every 4th frame anyway.

    Returns (frames, native_fps, n_frames_scanned) — the scan count is the number
    of source frames traversed, so callers can recover the real clip duration.
    """
    cap = cv2.VideoCapture(path)
    if not cap.isOpened():
        raise IOError(f"Could not open video: {path}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
    limit = int(max_seconds * fps) if (max_seconds and fps > 0) else None

    frames: list[np.ndarray] = []
    scanned = 0
    stride = max(1, int(stride))
    while len(frames) < max_frames:
        ok, frame = cap.read()
        if not ok:
            break
        if scanned % stride == 0:
            frames.append(frame)
        scanned += 1
        if limit is not None and scanned >= limit:
            break
    cap.release()

    if not frames:
        raise IOError(f"Video contained zero readable frames: {path}")
    return np.stack(frames, axis=0), float(fps), scanned


def write_preview_mp4(video: torch.Tensor, path: str, fps: float = 20.0) -> str:
    """Write a (T, C, H, W) uint8 RGB clip to `path` so users see the real model input.

    Uses imageio-ffmpeg (libx264) rather than cv2.VideoWriter: the pip OpenCV
    builds ship no H.264 encoder, and their mp4v fallback produces MPEG-4 Part 2,
    which most browsers refuse to play. A near-lossless CRF keeps the 1 px
    landmark overlay from being smeared away by compression.
    """
    import imageio.v2 as imageio

    frames = video.permute(0, 2, 3, 1).contiguous().numpy()  # (T, H, W, C) RGB
    writer = imageio.get_writer(
        path,
        fps=fps,
        codec="libx264",
        quality=None,
        output_params=["-crf", "12", "-pix_fmt", "yuv420p"],
        macro_block_size=1,   # dims are already multiples of 32; don't let it pad
    )
    try:
        for frame in frames:
            writer.append_data(frame)
    finally:
        writer.close()
    return path