src/fpgm/training/appearance_dataset.py

"""``torch.utils.data.Dataset`` over the appearance pairs, for the VACE fine-tune.

This is the second, much thinner training input alongside
:class:`~fpgm.training.dataset.WindowBundleDataset`, and the two exist for
different failure modes:

* ``WindowBundleDataset`` feeds the datagen bundles -- false-colour depth/seg/
  normal buffers composited into a control. Those are what the *deployed*
  generator will be conditioned on, and they need every lossless-decode
  guarantee ``bundle_io`` provides.
* This class feeds photograph-like controls: a real frame with a URDF-rendered
  Franka composited over the real arm (``scripts/appearance_control.py``). Its
  only job is to teach the model what a Franka *looks like*, which VACE-1.3B
  measurably does not know zero-shot -- handed a control whose robot was already
  a clean, correctly-posed white Franka + black Robotiq, it still drew a yellow
  toy arm from about frame 50.

Because the control here is an ordinary H.264 photograph, none of the lossless
machinery applies: there are no integer ids to corrupt, so a plain OpenCV decode
is not a shortcut, it is the whole contract.

The emitted dict has exactly the keys ``WanTrainingModule`` consumes, identical
to ``WindowBundleDataset.__getitem__`` -- ``video`` / ``prompt`` / ``vace_video``
/ ``vace_reference_image`` -- so either dataset can be handed to the same
trainer without a branch downstream.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image

from fpgm.types import DataError

#: VACE's training resolution. Both videos and the reference go through the same
#: cover-crop to it, so a control pixel and its target pixel stay the same pixel.
DEFAULT_RESOLUTION = (832, 480)

#: Wan's VAE has temporal stride 4 and encodes a leading frame on its own, so a
#: window length must be ``4k+1``. 81 is the length every datagen window already
#: uses; keeping it identical means a mixed run has one frame count, not two.
WINDOW_FRAMES = 81


class AppearancePairError(DataError):
    """A pair directory is unreadable or its two videos disagree."""


@dataclass(frozen=True)
class AppearanceWindow:
    """One ``(control, target)`` window: a pair directory plus a start frame."""

    pair_dir: Path
    start: int
    n_frames: int
    caption: str
    uuid: str
    camera_serial: str

    @property
    def name(self) -> str:
        return f"{self.uuid}/{self.camera_serial}/f{self.start:05d}"


def _cover_crop(image: np.ndarray, out_w: int, out_h: int) -> np.ndarray:
    """Scale-to-cover + centre-crop, the same transform ``export_vace._CoverCrop`` applies.

    Reimplemented rather than imported because that class is private to the
    datagen exporter and carries an intrinsics method this path has no camera
    for; the pixel arithmetic is deliberately identical so a mixed training run
    sees one framing convention, not two.
    """
    in_h, in_w = image.shape[:2]
    if (in_w, in_h) == (out_w, out_h):
        return image           # the pairs are written at the training resolution
    scale = max(out_w / in_w, out_h / in_h)
    rw, rh = int(round(in_w * scale)), int(round(in_h * scale))
    resized = cv2.resize(
        image, (rw, rh),
        interpolation=cv2.INTER_AREA if scale < 1.0 else cv2.INTER_LINEAR,
    )
    x0, y0 = int(round((rw - out_w) / 2.0)), int(round((rh - out_h) / 2.0))
    return resized[y0:y0 + out_h, x0:x0 + out_w]


def _read_window(path: Path, start: int, n: int, out_w: int, out_h: int) -> np.ndarray:
    """Decode ``n`` frames from ``start`` as cover-cropped RGB, ``(n, out_h, out_w, 3)``.

    Seeks with ``CAP_PROP_POS_FRAMES`` and then reads sequentially. The seek is
    only ever to a window start, so a keyframe-inexact seek costs at most a few
    frames of position error on a clip that was written at crf 14 with default
    GOP -- and both videos of a pair are seeked identically, so any such error is
    common-mode and does not desynchronise control from target.
    """
    cap = cv2.VideoCapture(str(path))
    if not cap.isOpened():
        raise AppearancePairError(f"could not open {path}")
    try:
        cap.set(cv2.CAP_PROP_POS_FRAMES, start)
        frames = []
        for _ in range(n):
            ok, bgr = cap.read()
            if not ok:
                break
            frames.append(_cover_crop(bgr, out_w, out_h)[:, :, ::-1])
    finally:
        cap.release()
    if len(frames) != n:
        raise AppearancePairError(
            f"{path}: wanted {n} frames from {start}, decoded {len(frames)}"
        )
    return np.stack(frames)


def discover_windows(
    root: Path,
    window_frames: int = WINDOW_FRAMES,
    stride: int | None = None,
    min_occlusion_ok: float = 0.25,
) -> list[AppearanceWindow]:
    """Every window of every pair directory under ``root``.

    Args:
        root: ``outputs/appearance_pairs`` -- one subdirectory per (episode,
            camera), each holding ``control.mp4``, ``target.mp4``, ``meta.json``.
        window_frames: frames per window; must stay ``4k+1``.
        stride: window start step. Defaults to ``window_frames`` (no overlap).
            A smaller stride multiplies the sample count from the same clips,
            at the cost of the extra windows being highly correlated.
        min_occlusion_ok: drop a clip whose ``occluded_px_fraction`` exceeds
            this. That number is the share of rendered robot pixels the depth
            test suppressed; a clip far above the ~1% typical value means the
            MoGe/render scale fit went wrong and the control has holes punched
            in the arm, which is the one defect that would actively teach the
            wrong appearance. This is a *diagnostic* filter on a measured
            quantity, not one of the datagen geometry gates -- those are about
            object pose and do not apply here.

    Raises:
        AppearancePairError: if ``root`` holds no usable pair.
    """
    stride = stride or window_frames
    if (window_frames - 1) % 4:
        raise ValueError(f"window_frames must be 4k+1, got {window_frames}")

    windows: list[AppearanceWindow] = []
    skipped: list[str] = []
    for meta_path in sorted(root.glob("*/meta.json")):
        meta = json.loads(meta_path.read_text())
        n = int(meta["n_frames"])
        if n < window_frames:
            skipped.append(f"{meta_path.parent.name}: {n} frames < {window_frames}")
            continue
        if float(meta.get("occluded_px_fraction", 0.0)) > min_occlusion_ok:
            skipped.append(
                f"{meta_path.parent.name}: occluded_px_fraction "
                f"{meta['occluded_px_fraction']:.3f} > {min_occlusion_ok}"
            )
            continue
        caption = (meta.get("caption") or "").strip()
        for start in range(0, n - window_frames + 1, stride):
            windows.append(AppearanceWindow(
                pair_dir=meta_path.parent, start=start, n_frames=window_frames,
                caption=caption, uuid=meta["uuid"],
                camera_serial=str(meta["camera_serial"]),
            ))
    if not windows:
        raise AppearancePairError(
            f"no usable appearance pairs under {root}"
            + (f" ({len(skipped)} skipped: {skipped[:3]})" if skipped else "")
        )
    return windows


class AppearancePairDataset(torch.utils.data.Dataset):
    """One item = one 81-frame ``(control, target)`` window.

    Batch size is effectively 1, matching ``WindowBundleDataset``: DiffSynth's
    training loop collates with ``lambda x: x[0]``.
    """

    def __init__(
        self,
        windows: list[AppearanceWindow],
        resolution: tuple[int, int] = DEFAULT_RESOLUTION,
        reference: str = "target_frame0",
    ) -> None:
        """
        Args:
            windows: from :func:`discover_windows`.
            resolution: ``(width, height)`` both videos are cover-cropped to.
            reference: what to hand VACE as ``vace_reference_image``.

                ``"target_frame0"`` -- the episode's real first frame, i.e. what
                the zero-shot probe used. It is decoded from ``target.mp4``
                rather than a stored PNG: a 0.9 MB per-clip copy of a frame that
                is already in the video is 25% of the clip's whole footprint for
                nothing. The real robot is *in* that frame, so
                the model can learn to copy its appearance from the reference
                rather than into its own weights. That is a real train/deploy
                mismatch: at datagen inference time the reference is
                ``ref_plate.png``, a background plate with no robot in it.

                ``"window_frame0"`` -- the window's own first *control* frame
                (CAD robot, real background), which removes the real robot from
                the reference entirely and forces the appearance into the
                weights, at the cost of no longer matching the probe.

            Neither is silently "correct"; the mismatch above is the thing to
            measure, so both are available and the choice is recorded per item
            in ``_reference``.
        """
        if not windows:
            raise AppearancePairError("AppearancePairDataset built with zero windows")
        if reference not in ("target_frame0", "window_frame0"):
            raise ValueError(f"unknown reference mode {reference!r}")
        self.windows = list(windows)
        self.resolution = resolution
        self.reference = reference

    def __len__(self) -> int:
        return len(self.windows)

    def __getitem__(self, index: int) -> dict:
        win = self.windows[index % len(self.windows)]
        out_w, out_h = self.resolution
        target = _read_window(win.pair_dir / "target.mp4", win.start, win.n_frames,
                              out_w, out_h)
        control = _read_window(win.pair_dir / "control.mp4", win.start, win.n_frames,
                               out_w, out_h)

        if self.reference == "window_frame0":
            ref = control[0]
        else:
            ref = _read_window(win.pair_dir / "target.mp4", 0, 1, out_w, out_h)[0]

        return {
            "video": [Image.fromarray(f) for f in target],
            "prompt": win.caption,
            "vace_video": [Image.fromarray(f) for f in control],
            "vace_reference_image": [Image.fromarray(np.ascontiguousarray(ref))],
            "_window_name": win.name,
            "_reference": self.reference,
        }