src/fpgm/training/bundle_io.py

"""Lossless decode + control-video compositing for one ``vace/window_*`` bundle.

**Why losslessness matters here, restated at the point it is actually enforced**
(not just documented upstream): ``control_seg.mkv`` stores small integer category
ids (0=bg, 1=robot, 2+=objects) as raw pixel values. Any lossy re-encode between
export and here -- or a lossy decode path that happens to produce the same
values 99% of the time -- would let ringing shift an id by +-1 near an edge,
silently relabelling a robot pixel as an object pixel. :func:`read_id_video`
therefore decodes via OpenCV (confirmed by direct test against the FFV1 file on
disk to round-trip exact integer values, see this module's own smoke check in
``tests/test_training_dataset.py``) and asserts every decoded pixel is in
``[0, max_id]`` rather than silently clipping.

``cv2``/``numpy`` are imported at module scope, matching the rest of this
codebase's convention for CPU-only libraries (e.g. ``fpgm.datagen.export_vace``
imports ``cv2`` at module scope too) -- only torch/DiffSynth are deferred to
call time elsewhere in this package, since those are the GPU-coupled imports
the project's "lazy heavy imports" rule is actually about.
"""

from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np

from fpgm.geometry.normals import decode_normals_rgb
from fpgm.training.types import BundleAssemblyConfig, WindowSample
from fpgm.types import DataError


class BundleShapeError(DataError):
    """A decoded video's frame count/resolution/dtype doesn't match ``sample.json``."""


def _read_video_bgr(path: Path, n_expected_frames: int) -> np.ndarray:
    """Decode every frame of ``path`` as BGR uint8, shape ``(T, H, W, 3)``."""
    cap = cv2.VideoCapture(str(path))
    if not cap.isOpened():
        raise BundleShapeError(f"could not open {path}")
    frames = []
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            frames.append(frame)
    finally:
        cap.release()
    if len(frames) != n_expected_frames:
        raise BundleShapeError(
            f"{path}: decoded {len(frames)} frames, sample.json declares {n_expected_frames}"
        )
    return np.stack(frames, axis=0)


def read_rgb_video(path: Path, n_expected_frames: int) -> np.ndarray:
    """Decode ``path`` as RGB uint8, shape ``(T, H, W, 3)``."""
    bgr = _read_video_bgr(path, n_expected_frames)
    return bgr[..., ::-1]


def read_gray_video(path: Path, n_expected_frames: int) -> np.ndarray:
    """Decode a grayscale-stored-as-3-channel video, shape ``(T, H, W)``.

    Used for ``control_depth.mkv`` (inverse-depth 8-bit, channel-replicated).
    Asserts the three channels agree exactly -- disagreement would mean the
    file isn't actually grayscale-replicated (a format assumption violation,
    not a value to silently average away).
    """
    bgr = _read_video_bgr(path, n_expected_frames)
    if not (np.array_equal(bgr[..., 0], bgr[..., 1]) and np.array_equal(bgr[..., 1], bgr[..., 2])):
        raise BundleShapeError(f"{path}: expected a channel-replicated grayscale video")
    return bgr[..., 0]


def read_id_video(path: Path, n_expected_frames: int, max_id: int) -> np.ndarray:
    """Decode ``control_seg.mkv`` losslessly, shape ``(T, H, W)`` uint8 ids.

    See module docstring: this is the one channel where a lossy decode would
    be a silent correctness bug, not just a quality regression.
    """
    ids = read_gray_video(path, n_expected_frames)
    bad = ids[ids > max_id]
    if bad.size > 0:
        raise BundleShapeError(
            f"{path}: found seg id(s) up to {int(ids.max())}, exceeding max_id={max_id} "
            "-- either the export uses more categories than BundleAssemblyConfig.max_seg_id "
            "expects, or this file is not actually a lossless id map."
        )
    return ids


def read_caption(path: Path) -> str:
    return path.read_text().strip()


def letterbox_to_white(
    image_rgb: np.ndarray, target_w: int, target_h: int, fill: int = 255
) -> np.ndarray:
    """Scale ``image_rgb`` to fit inside ``(target_w, target_h)`` and pad with ``fill``.

    Reproduces the established Wan-VACE reference-image convention (letterbox
    on a white canvas, *not* scale-to-cover + crop -- that transform is for
    control videos only). Keeping the Dataset's own reference-image handling
    consistent with that convention matters because DiffSynth's
    ``preprocess_video``/``preprocess_image`` (``diffsynth/diffusion/base_pipeline.py``)
    does **no** resizing at all -- it requires every PIL image already be the
    target (H, W), so whichever convention this function picks *is* the
    convention the model trains on. Matching the inference-time letterbox
    (rather than e.g. center-cropping the plate) avoids a second train/inference
    mismatch on top of the depth-format one this whole finetune exists to fix.
    """
    h, w = image_rgb.shape[:2]
    scale = min(target_w / w, target_h / h)
    new_w, new_h = max(1, round(w * scale)), max(1, round(h * scale))
    resized = cv2.resize(image_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA)
    canvas = np.full((target_h, target_w, 3), fill, dtype=np.uint8)
    y0 = (target_h - new_h) // 2
    x0 = (target_w - new_w) // 2
    canvas[y0 : y0 + new_h, x0 : x0 + new_w] = resized
    return canvas


def normal_rgb_to_z_channel(normal_rgb: np.ndarray) -> np.ndarray:
    """``(T, H, W, 3)`` encoded normal maps -> ``(T, H, W)`` uint8 Z-channel.

    Reuses :func:`fpgm.geometry.normals.decode_normals_rgb` (the exact inverse
    of the encoder ``fpgm.datagen.export_vace`` used to write
    ``control_normal.mp4``) rather than re-deriving the RGB<->[-1,1] mapping
    here, so a future change to that convention can't silently desync the two
    call sites. Applied per-frame: :func:`decode_normals_rgb` only accepts a
    single ``(H, W, 3)`` frame, not a batched ``(T, H, W, 3)`` video.
    """
    z_frames = [decode_normals_rgb(normal_rgb[t])[..., 2] for t in range(normal_rgb.shape[0])]
    z = np.stack(z_frames, axis=0)
    return np.clip((z * 0.5 + 0.5) * 255.0, 0.0, 255.0).astype(np.uint8)


def compose_control_rgb(
    depth_gray: np.ndarray,
    seg_ids: np.ndarray,
    normal_z: np.ndarray,
    max_seg_id: int,
) -> np.ndarray:
    """Composite the three lossless/near-lossless control channels into one RGB video.

    ``(T, H, W)`` depth/seg/normal-Z -> ``(T, H, W, 3)`` uint8, R=depth,
    G=seg id normalised to ``[0, 255]`` by ``max_seg_id``, B=normal-Z. See
    :class:`~fpgm.training.types.BundleAssemblyConfig`'s docstring for why a
    single composited video (not three separate control inputs) is what the
    trainer's VACE branch actually consumes.
    """
    if not (depth_gray.shape == seg_ids.shape == normal_z.shape):
        raise BundleShapeError(
            f"channel shape mismatch: depth={depth_gray.shape} "
            f"seg={seg_ids.shape} normal={normal_z.shape}"
        )
    seg_scaled = seg_ids.astype(np.float32) * (255.0 / max_seg_id)
    seg_scaled = np.clip(np.round(seg_scaled), 0, 255).astype(np.uint8)
    return np.stack([depth_gray, seg_scaled, normal_z], axis=-1)


def load_window_arrays(
    sample: WindowSample, cfg: BundleAssemblyConfig
) -> dict[str, np.ndarray | str]:
    """Decode every file in ``sample`` into the raw numpy arrays the Dataset needs.

    Returns a dict with keys ``target`` ``(T,H,W,3)`` uint8, ``control``
    ``(T,H,W,3)`` uint8 (the composite), ``fg_mask`` ``(T,H,W)`` bool,
    ``ref_plate`` ``(H,W,3)`` uint8 (native resolution, *not yet* letterboxed
    -- letterboxing happens in the Dataset after augmentation config decides
    the target size), and ``caption`` str. No augmentation or tensor
    conversion happens here -- this function is pure decode + composite, kept
    separate so :mod:`fpgm.training.augment` can be unit-tested against
    hand-built arrays without touching the filesystem.
    """
    n = sample.frame_range[1] - sample.frame_range[0]
    depth = read_gray_video(sample.control_depth, n)
    seg = read_id_video(sample.control_seg, n, cfg.max_seg_id)
    normal_rgb = read_rgb_video(sample.control_normal, n)
    normal_z = normal_rgb_to_z_channel(normal_rgb)
    control = compose_control_rgb(depth, seg, normal_z, cfg.max_seg_id)

    target = read_rgb_video(sample.target_video, n)
    fg_mask_rgb = read_rgb_video(sample.fg_mask, n)
    fg_mask = fg_mask_rgb[..., 0] > 127

    ref_plate_bgr = cv2.imread(str(sample.ref_plate), cv2.IMREAD_COLOR)
    if ref_plate_bgr is None:
        raise BundleShapeError(f"could not read {sample.ref_plate}")
    ref_plate = ref_plate_bgr[..., ::-1]

    caption = read_caption(sample.caption_path)

    return {
        "target": target,
        "control": control,
        "fg_mask": fg_mask,
        "ref_plate": ref_plate,
        "caption": caption,
    }