src/fpgm/training/augment.py

"""Train<->inference-gap augmentations: pose noise, mask dilate/erode, channel dropout.

Every function here is pure numpy/OpenCV (no torch), takes an explicit
``numpy.random.Generator``, and is independently toggleable -- see
:class:`fpgm.training.types.AugmentConfig`. :func:`apply_augmentations` is the
single entry point :mod:`fpgm.training.dataset` calls; the three primitives
below are exported separately so each can be unit-tested in isolation against
hand-built synthetic arrays (no filesystem, no GPU).
"""

from __future__ import annotations

import cv2
import numpy as np

from fpgm.training.types import AugmentConfig


def pose_noise(
    control: np.ndarray,
    fg_mask: np.ndarray,
    rng: np.random.Generator,
    translate_px_std: float,
    rotate_deg_std: float,
) -> np.ndarray:
    """Rigid-jitter the foreground region of the composited control video only.

    **Why a 2D affine warp of the raster, not a 3D pose perturbation re-rendered
    through meshes.** True pose noise (perturb ``T_world_obj`` by +-2mm/+-1deg,
    re-render depth/seg/normal) needs the mesh + camera-intrinsics + pyrender
    machinery that ``fpgm.datagen.export_vace`` (a sibling agent's owned
    module, not duplicated here) already owns, and pulling that dependency
    into a lightweight Dataset augmentation would mean re-instantiating an
    OSMesa-backed GPU/CPU renderer per training sample -- disproportionate to
    what this augmentation needs to achieve. A small rigid affine warp of the
    already-rendered foreground pixels is a raster-level approximation of the
    same effect (an object rendered from a slightly-off pose looks, to first
    order, like the correctly-rendered object translated/rotated a few
    pixels) -- it will not reproduce genuine parallax/occlusion changes a
    real re-render would, but its purpose is only to keep the model from
    overfitting to pixel-exact conditioning, which this achieves.

    One random affine draw per call, applied identically to every frame in
    the window (not redrawn per frame) -- a per-frame-random jitter would
    look like flicker no real pose-estimation error produces; a real pose
    bias is temporally coherent within a short window.

    Args:
        control: ``(T, H, W, 3)`` uint8 composited control video.
        fg_mask: ``(T, H, W)`` bool, the region allowed to move (background
            pixels are left untouched regardless of the warp).
        rng: source of randomness.
        translate_px_std: stddev of the (independent x/y) pixel translation.
        rotate_deg_std: stddev of the in-plane rotation, degrees.

    Returns:
        A new ``(T, H, W, 3)`` uint8 array; ``control`` is not modified in place.
    """
    if control.shape[:3] != fg_mask.shape:
        raise ValueError(
            f"control {control.shape[:3]} and fg_mask {fg_mask.shape} disagree on (T,H,W)"
        )
    t_frames, h, w = fg_mask.shape
    dx = float(rng.normal(0.0, translate_px_std))
    dy = float(rng.normal(0.0, translate_px_std))
    angle_deg = float(rng.normal(0.0, rotate_deg_std))

    center = (w / 2.0, h / 2.0)
    matrix = cv2.getRotationMatrix2D(center, angle_deg, 1.0)
    matrix[0, 2] += dx
    matrix[1, 2] += dy

    out = control.copy()
    for t in range(t_frames):
        if not fg_mask[t].any():
            continue
        warped_frame = cv2.warpAffine(
            control[t], matrix, (w, h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE
        )
        warped_mask = cv2.warpAffine(
            fg_mask[t].astype(np.uint8) * 255, matrix, (w, h),
            flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0,
        ) > 127
        combined = fg_mask[t] | warped_mask  # cover both the vacated and the arrived-at footprint
        out[t] = np.where(combined[..., None], warped_frame, control[t])
    return out


def mask_dilate_erode(
    fg_mask: np.ndarray, rng: np.random.Generator, max_kernel_px: int
) -> np.ndarray:
    """Randomly dilate, erode, or leave ``fg_mask`` unchanged.

    One draw per call (op choice + kernel size), applied identically to every
    frame -- same temporal-coherence argument as :func:`pose_noise`: a mask
    generator's systematic bias (consistently a few px too tight/loose) looks
    like a fixed dilate/erode, not per-frame flicker.

    Args:
        fg_mask: ``(T, H, W)`` bool.
        rng: source of randomness.
        max_kernel_px: inclusive upper bound on the structuring element's
            half-width; ``0`` degenerates to always-unchanged.

    Returns:
        A new ``(T, H, W)`` bool array.
    """
    if max_kernel_px <= 0:
        return fg_mask.copy()
    op = rng.choice(("dilate", "erode", "none"))
    if op == "none":
        return fg_mask.copy()
    kernel_px = int(rng.integers(1, max_kernel_px + 1))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * kernel_px + 1, 2 * kernel_px + 1))
    fn = cv2.dilate if op == "dilate" else cv2.erode
    out = np.empty_like(fg_mask)
    for t in range(fg_mask.shape[0]):
        out[t] = fn(fg_mask[t].astype(np.uint8), kernel) > 0
    return out


def channel_dropout(control: np.ndarray, rng: np.random.Generator, prob: float) -> np.ndarray:
    """Independently zero each of the 3 composited channels with probability ``prob``.

    Whole-window, not per-frame (same temporal-coherence reasoning as the
    other two augmentations): this simulates "this render channel was
    unavailable/degenerate for this clip", not a per-frame glitch.

    Args:
        control: ``(T, H, W, 3)`` uint8.
        rng: source of randomness.
        prob: independent per-channel dropout probability.

    Returns:
        A new ``(T, H, W, 3)`` uint8 array.
    """
    out = control.copy()
    for c in range(control.shape[-1]):
        if rng.random() < prob:
            out[..., c] = 0
    return out


def apply_augmentations(
    control: np.ndarray, fg_mask: np.ndarray, cfg: AugmentConfig, sample_index: int
) -> tuple[np.ndarray, np.ndarray]:
    """Apply the enabled augmentations in :class:`~fpgm.training.types.AugmentConfig` order.

    Order: mask dilate/erode first (its output is both returned and used as
    the foreground region for pose noise, so a systematically-off mask and a
    pose jitter compose the way they would in reality), then pose noise, then
    channel dropout last (it should be able to zero a channel pose noise just
    perturbed).

    ``target``/``ref_plate``/``caption`` are never touched here: the real
    video is ground truth and the reference plate is a real photograph --
    both must reach the model unmodified for the loss/context to mean
    anything (only the *rendered control conditioning* is uncertain enough
    at inference time to warrant jitter).

    Args:
        control: ``(T, H, W, 3)`` uint8 composited control video.
        fg_mask: ``(T, H, W)`` bool.
        cfg: which augmentations are enabled and their magnitudes.
        sample_index: mixed with ``cfg.seed`` to derive this call's RNG, so a
            given (seed, index) is reproducible across runs/workers while
            different indices/epochs draw independently.

    Returns:
        ``(control_aug, fg_mask_aug)``.
    """
    rng = np.random.default_rng((cfg.seed, sample_index))

    mask_aug = fg_mask
    if cfg.enable_mask_dilate_erode:
        mask_aug = mask_dilate_erode(mask_aug, rng, cfg.mask_dilate_erode_max_px)

    control_aug = control
    if cfg.enable_pose_noise:
        translate_px_std = cfg.pose_noise_translate_mm_std * cfg.pose_noise_px_per_mm
        control_aug = pose_noise(
            control_aug, mask_aug, rng, translate_px_std, cfg.pose_noise_rotate_deg_std
        )

    if cfg.enable_channel_dropout:
        control_aug = channel_dropout(control_aug, rng, cfg.channel_dropout_prob)

    return control_aug, mask_aug