src/fpgm/pipeline/frames.py

"""Clip frame extraction.

A ``*_flows.h5`` clip covers a sub-range ``{start}:{end}`` of an episode, while the
mp4 holds the whole episode. Materialising exactly the clip's frames as a numbered
JPEG folder buys two things:

* SAM 3.1 accepts a ``<index>.jpg`` folder directly, and both it and TAPNext++ then
  see the *same* frames in the *same* order, so index 0 means the same instant to
  the segmenter, the tracker and ``scene_flows[0]``.
* Index alignment stops being an assumption and becomes a construction, which
  removes the largest class of silent off-by-one errors in this pipeline.

``stride`` exists because that "same instant" mapping is not always 1:1: PointWorld
annotations run at half the video/trajectory frame rate (see
:class:`fpgm.types.ClipTiming`'s docstring), so ``[start_frame, end_frame)`` alone
would materialise consecutive video frames for what are actually every-other-frame
annotations. ``stride`` selects every Nth frame in that range instead, so callers
pass the *mapped* video range (via ``ClipTiming.clip_frame_to_video_frame``) and the
matching stride, and get back exactly one extracted frame per annotated clip frame.
"""

from __future__ import annotations

import shutil
import tempfile
from collections.abc import Iterator
from pathlib import Path

import cv2
import numpy as np

from fpgm.utils.logging import get_logger

logger = get_logger(__name__)


class ClipFrameSource:
    """Extracts one clip's frame range from an mp4 into a numbered JPEG folder.

    Use as a context manager so the extracted frames are cleaned up even when a
    clip fails partway through -- these are a few hundred MB per clip on a disk
    that is already 98% full.
    """

    def __init__(
        self,
        video_path: str | Path,
        start_frame: int,
        end_frame: int,
        work_dir: str | Path | None = None,
        jpeg_quality: int = 95,
        stride: int = 1,
    ) -> None:
        """See the class/module docstrings for ``stride``.

        Args:
            video_path: Source mp4.
            start_frame: First video frame index to extract (inclusive).
            end_frame: Video frame index to stop at (exclusive) -- i.e. extraction
                covers ``range(start_frame, end_frame, stride)``.
            work_dir: Directory to write JPEGs into; a temp dir is created and
                owned (cleaned up on ``cleanup()``/context exit) if omitted.
            jpeg_quality: JPEG quality passed to ``cv2.imwrite``.
            stride: Extract every ``stride``-th frame starting at ``start_frame``,
                rather than every consecutive one. Defaults to 1 (the original,
                consecutive-frame behaviour), which keeps every existing caller
                (e.g. the robot renderer, which genuinely wants every video frame)
                unaffected.

        Raises:
            ValueError: If ``stride < 1``.
        """
        if stride < 1:
            raise ValueError(f"stride must be >= 1, got {stride}")
        self.video_path = Path(video_path)
        self.start_frame = int(start_frame)
        self.end_frame = int(end_frame)
        self.jpeg_quality = int(jpeg_quality)
        self.stride = int(stride)
        self._work_dir = Path(work_dir) if work_dir else None
        self._owns_dir = work_dir is None
        self.frame_dir: Path | None = None
        self._n_frames = 0
        self._size: tuple[int, int] | None = None

    @property
    def n_frames(self) -> int:
        return self._n_frames

    @property
    def resolution(self) -> tuple[int, int]:
        """(width, height) of the extracted frames."""
        if self._size is None:
            raise RuntimeError("frames have not been extracted yet")
        return self._size

    def extract(self) -> Path:
        """Write frames ``range(start_frame, end_frame, stride)`` as ``00000.jpg`` onwards."""
        if self._work_dir is None:
            self._work_dir = Path(tempfile.mkdtemp(prefix="fpgm_frames_"))
        self.frame_dir = self._work_dir
        self.frame_dir.mkdir(parents=True, exist_ok=True)

        cap = cv2.VideoCapture(str(self.video_path))
        if not cap.isOpened():
            raise FileNotFoundError(f"could not open video: {self.video_path}")

        try:
            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            end = min(self.end_frame, total) if total > 0 else self.end_frame
            if end <= self.start_frame:
                raise ValueError(
                    f"empty clip range [{self.start_frame}, {end}) for {self.video_path} "
                    f"(video reports {total} frames)"
                )
            if total > 0 and self.end_frame > total:
                # Not fatal, but the caller's clip metadata disagrees with the media.
                logger.warning(
                    "clip end %d exceeds video length %d for %s; truncating",
                    self.end_frame,
                    total,
                    self.video_path.name,
                )

            # Seeking is unreliable on some codecs, so read sequentially and skip.
            # At ~130 frames per episode this costs nothing and is always correct.
            written = 0
            for idx in range(end):
                ok, frame = cap.read()
                if not ok:
                    break
                if idx < self.start_frame:
                    continue
                if (idx - self.start_frame) % self.stride != 0:
                    continue
                out_path = self.frame_dir / f"{written:05d}.jpg"
                cv2.imwrite(
                    str(out_path),
                    frame,
                    [int(cv2.IMWRITE_JPEG_QUALITY), self.jpeg_quality],
                )
                if self._size is None:
                    self._size = (frame.shape[1], frame.shape[0])
                written += 1
        finally:
            cap.release()

        if written == 0:
            raise ValueError(f"no frames extracted from {self.video_path}")

        self._n_frames = written
        last_frame = self.start_frame + (written - 1) * self.stride
        logger.info(
            "extracted %d frames [%d, %d] stride=%d from %s -> %s",
            written,
            self.start_frame,
            last_frame,
            self.stride,
            self.video_path.name,
            self.frame_dir,
        )
        return self.frame_dir

    def iter_rgb(self) -> Iterator[np.ndarray]:
        """Yield extracted frames as RGB uint8 arrays, lazily."""
        for path in self.paths():
            bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
            if bgr is None:
                raise RuntimeError(f"could not read extracted frame {path}")
            yield cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)

    def iter_bgr(self) -> Iterator[np.ndarray]:
        """Yield extracted frames as BGR uint8 arrays, lazily."""
        for path in self.paths():
            bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
            if bgr is None:
                raise RuntimeError(f"could not read extracted frame {path}")
            yield bgr

    def paths(self) -> list[Path]:
        if self.frame_dir is None:
            raise RuntimeError("call extract() first")
        return sorted(self.frame_dir.glob("*.jpg"), key=lambda p: int(p.stem))

    def read(self, index: int) -> np.ndarray:
        """Read a single extracted frame as BGR."""
        paths = self.paths()
        bgr = cv2.imread(str(paths[index]), cv2.IMREAD_COLOR)
        if bgr is None:
            raise RuntimeError(f"could not read extracted frame {paths[index]}")
        return bgr

    def cleanup(self) -> None:
        if self._owns_dir and self.frame_dir is not None and self.frame_dir.exists():
            shutil.rmtree(self.frame_dir, ignore_errors=True)
            self.frame_dir = None

    def __enter__(self) -> "ClipFrameSource":
        self.extract()
        return self

    def __exit__(self, *exc_info: object) -> None:
        self.cleanup()