src/fpgm/datagen/plate.py

"""Temporal-median background plate: a real photographic ``ref_plate.png``.

Promoted (reimplemented cleanly, not imported -- see the module docstring of
``scripts/export_cosmos_input.py``, which owns the working reference this is
based on and is left untouched) as its own stage because S8's VACE export needs
exactly this artifact, and it is cheap enough to build once per episode/camera
independent of every other stage.

**This only works because the capture camera is static to machine precision on
this dataset** (verified in earlier work on this project, not re-derived here).
A per-pixel temporal median over every real video frame removes anything that
*moves* against the background -- principally the robot arm, which occupies any
given background pixel for well under half of a typical episode -- because the
median of a pixel's whole time series lands on whichever value it took most of
the time, i.e. the background. If the camera itself moved (panned, zoomed, or
was bumped), this same operation would instead blend genuinely *different*
background content -- different physical points in the scene passing under the
same pixel coordinate over time -- into one blurred, physically meaningless
average that looks like a plate but is not one. :func:`compute_background_plate`
therefore always runs :func:`_check_static_camera` and warns loudly
(:class:`NonStaticCameraWarning`) rather than silently returning a plate that
would be wrong in a way no downstream consumer could detect from the image
alone.
"""

from __future__ import annotations

import warnings
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np

from fpgm.datagen.cache import StageCache
from fpgm.types import DataError
from fpgm.utils.io import ensure_dir
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: Corner-patch side length as a fraction of (height, width). A corner, not the
#: frame centre, is checked deliberately: the robot's workspace (and therefore
#: essentially all genuine foreground motion) is framed centrally in this
#: dataset's captures, so a corner patch is the region least likely to ever be
#: legitimately occluded by the moving arm/object -- anything that moves there
#: is either camera motion or a lighting change, both of which violate this
#: plate's static-background assumption.
_CORNER_FRAC = 0.12

#: Median abs frame-to-frame pixel difference (0-255 scale) in the corner patch
#: above which the camera is flagged as probably not static. Ordinary sensor
#: noise on a genuinely static patch is a fraction of one 8-bit level once
#: median-pooled over many frame-pairs; a panning/zooming camera sweeps entirely
#: different scene content through the corner and pushes this well into the
#: tens.
_STATIC_CAMERA_WARN_THRESHOLD = 6.0


class NonStaticCameraWarning(UserWarning):
    """The corner-patch check found more frame-to-frame motion than a static
    camera should ever produce -- the temporal-median plate this run produced is
    likely a physically meaningless blur of different background content, not a
    clean photograph. See the module docstring.
    """


@dataclass(frozen=True)
class PlateStats:
    """Diagnostics from one :func:`compute_background_plate` run."""

    n_frames: int
    corner_median_abs_diff: float  # 0-255 scale; see _STATIC_CAMERA_WARN_THRESHOLD.
    static_camera_ok: bool


def _check_static_camera(
    frames_bgr: np.ndarray, corner_frac: float, threshold: float
) -> PlateStats:
    """Median abs frame-to-frame difference in a corner patch. See module docstring.

    Args:
        frames_bgr: ``(T, H, W, 3)`` uint8, every decoded video frame.
        corner_frac: Patch side length as a fraction of ``(H, W)``.
        threshold: Above this, :attr:`PlateStats.static_camera_ok` is ``False``.
    """
    n, h, w = frames_bgr.shape[:3]
    patch_h = max(1, int(round(h * corner_frac)))
    patch_w = max(1, int(round(w * corner_frac)))
    patch = frames_bgr[:, :patch_h, :patch_w, :].astype(np.float32)
    diffs = np.abs(np.diff(patch, axis=0))
    median_diff = float(np.median(diffs)) if diffs.size else 0.0
    return PlateStats(
        n_frames=n, corner_median_abs_diff=median_diff, static_camera_ok=median_diff <= threshold
    )


def compute_background_plate(
    video_path: str | Path,
    *,
    corner_frac: float = _CORNER_FRAC,
    static_camera_threshold: float = _STATIC_CAMERA_WARN_THRESHOLD,
) -> tuple[np.ndarray, PlateStats]:
    """Per-pixel temporal median of a whole video -> a robot-free RGB plate.

    Args:
        video_path: Path to the real episode video (e.g. the raw DROID mp4).
        corner_frac: See :data:`_CORNER_FRAC`.
        static_camera_threshold: See :data:`_STATIC_CAMERA_WARN_THRESHOLD`.

    Returns:
        ``(plate_rgb, stats)`` -- ``plate_rgb`` is ``(H, W, 3)`` uint8 at the
        video's own resolution (no resize: the caller decides what resolution
        it needs the plate at).

    Raises:
        DataError: if no frames could be decoded from ``video_path``.

    Warns:
        NonStaticCameraWarning: if the corner-patch check fails -- the plate is
            still returned (a best-effort blur is more useful for a human to
            eyeball-diagnose than nothing), but callers must not trust it as a
            clean background without investigating.
    """
    import cv2

    cap = cv2.VideoCapture(str(video_path))
    frames: list[np.ndarray] = []
    try:
        while True:
            ok, bgr = cap.read()
            if not ok:
                break
            frames.append(bgr)
    finally:
        cap.release()
    if not frames:
        raise DataError(f"no frames decoded from {video_path}")

    stacked = np.stack(frames, axis=0)  # (T, H, W, 3) BGR uint8
    plate_bgr = np.median(stacked, axis=0).astype(np.uint8)

    stats = _check_static_camera(stacked, corner_frac, static_camera_threshold)
    if not stats.static_camera_ok:
        message = (
            f"{video_path}: corner-patch median abs frame-to-frame difference "
            f"{stats.corner_median_abs_diff:.2f} exceeds the static-camera threshold "
            f"{static_camera_threshold:.2f} (over {stats.n_frames} frames). This plate "
            "assumes a fixed camera pose -- temporal median only removes a *moving "
            "foreground* against a *static background*. If the camera panned, zoomed, "
            "or was bumped during capture, the median instead blends genuinely "
            "different background content into one blurred, physically meaningless "
            "image with no visual sign anything went wrong."
        )
        logger.warning(message)
        warnings.warn(message, NonStaticCameraWarning, stacklevel=2)

    plate_rgb = cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2RGB)
    return plate_rgb, stats


class PlateStage:
    """Orchestration: build once, cache via :class:`StageCache`, write ``plate.png``."""

    def run(self, *, video_path: str | Path, cache: StageCache) -> tuple[np.ndarray, PlateStats]:
        """Build (or reuse) the background plate for one episode/camera.

        Args:
            video_path: Path to the real episode video.
            cache: Rooted at ``outputs/datagen/<uuid>/<camera_serial>/`` -- the
                plate is written to ``<cache.root>/master/plate.png`` (see
                :class:`~fpgm.datagen.dense_depth.DenseDepthStage` for why the
                deliverable path is decoupled from the cache's own bookkeeping
                directory).

        Returns:
            ``(plate_rgb, stats)``.
        """
        import cv2

        stage_name = "plate"
        master_dir = ensure_dir(cache.root / "master")
        plate_path = master_dir / "plate.png"
        video_path = Path(video_path)

        fingerprint = {
            "video_path": str(video_path),
            "video_mtime": video_path.stat().st_mtime,
            "video_size": video_path.stat().st_size,
            "code_version": "plate.v1",
        }

        if cache.is_fresh(stage_name, fingerprint) and plate_path.exists():
            logger.info("plate: reusing cached %s", plate_path)
            plate_bgr = cv2.imread(str(plate_path))
            meta = cache.read_meta(stage_name) or {}
            stats = PlateStats(**meta.get("payload", {})["stats"])
            return cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2RGB), stats

        plate_rgb, stats = compute_background_plate(video_path)
        cv2.imwrite(str(plate_path), cv2.cvtColor(plate_rgb, cv2.COLOR_RGB2BGR))

        limitations = [
            "Valid only because the capture camera is static to machine precision on "
            "this dataset; a panned/zoomed/bumped recording would blend genuinely "
            "different background content into one meaningless blur -- see "
            "PlateStats.static_camera_ok / corner_median_abs_diff.",
        ]
        if not stats.static_camera_ok:
            limitations.append(
                f"static-camera check FAILED this run (corner_median_abs_diff="
                f"{stats.corner_median_abs_diff:.2f} > threshold); plate.png is likely "
                "not a clean background photograph."
            )

        cache.write_meta(
            stage_name,
            fingerprint,
            payload={"stats": asdict(stats), "plate_png": str(plate_path)},
            limitations=limitations,
        )
        return plate_rgb, stats