src/fpgm/objects/crop.py

"""Stage 1: SAM 3.1 mask -> tight RGBA crop of the manipulated object.

The rest of the pipeline (mesh generation, alignment) only ever sees the crop
this stage produces, so the two judgment calls made here -- *which frame* and
*how tight a box* -- are the highest-leverage decisions in the whole pipeline:
a bad choice here cannot be corrected downstream.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import cv2
import numpy as np

from fpgm.objects.types import ObjectCrop, StageArtifacts
from fpgm.utils.io import ensure_dir
from fpgm.utils.logging import get_logger
from fpgm.viz.overlays import color_for, overlay_mask

logger = get_logger(__name__)

_SELECT_METHODS = ("largest", "least_truncated")
_INPAINT_METHODS = {"telea": cv2.INPAINT_TELEA, "ns": cv2.INPAINT_NS}

#: Above this occluded fraction, cv2.inpaint's PDE fill is smearing neighbouring
#: texture rather than reconstructing anything -- advisory only, not enforced,
#: see :func:`inpaint_occlusion`.
OCCLUSION_TRUST_THRESHOLD = 0.3


def crop_from_mask(
    frame_bgr: np.ndarray,
    mask: np.ndarray,
    padding: float = 0.15,
    min_size: int = 64,
    square: bool = True,
    frame_idx: int = 0,
    prompt: str = "",
) -> ObjectCrop:
    """Tight, padded RGBA crop of ``mask`` from ``frame_bgr``.

    The alpha channel is set to the mask so an image-to-3D generator sees a
    clean cutout rather than background clutter it would have to segment
    itself. ``square`` defaults on because these models are almost universally
    trained on roughly-square inputs; a wide/tall box would be silently
    letterboxed or distorted downstream.

    Args:
        frame_bgr: ``(H, W, 3)`` uint8 BGR frame.
        mask: ``(H, W)`` boolean object mask, same resolution as ``frame_bgr``.
        padding: Fractional padding added to each side of the tight bbox,
            relative to the bbox's own width/height.
        min_size: Minimum crop side length in pixels, enforced after padding.
        square: Expand the shorter side to match the longer one.
        frame_idx: Video frame index this crop came from, carried through for
            provenance/debug filenames.
        prompt: Text prompt that produced ``mask``, carried through for
            provenance.

    Returns:
        The cropped :class:`~fpgm.objects.types.ObjectCrop`.

    Raises:
        ValueError: If ``mask`` is empty or its shape disagrees with
            ``frame_bgr``.
    """
    mask = np.asarray(mask, dtype=bool)
    if frame_bgr.shape[:2] != mask.shape:
        raise ValueError(
            f"crop_from_mask: frame shape {frame_bgr.shape[:2]} != mask shape {mask.shape}"
        )
    if not mask.any():
        raise ValueError("crop_from_mask: mask is empty")

    h, w = mask.shape
    x0, y0, x1, y1 = _padded_bbox(mask, padding, min_size, square, w, h)

    crop_bgr = frame_bgr[y0:y1, x0:x1]
    crop_mask = mask[y0:y1, x0:x1]
    rgba = np.empty((*crop_mask.shape, 4), dtype=np.uint8)
    rgba[..., :3] = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2RGB)
    rgba[..., 3] = crop_mask.astype(np.uint8) * 255

    return ObjectCrop(
        frame_idx=frame_idx,
        rgba=rgba,
        bbox_xyxy=(x0, y0, x1, y1),
        mask_full=mask,
        prompt=prompt,
    )


def _padded_bbox(
    mask: np.ndarray, padding: float, min_size: int, square: bool, w: int, h: int
) -> tuple[int, int, int, int]:
    """Tight bbox of ``mask``, padded / squared / min-sized, clipped to the frame.

    Clipping shifts the box to fit inside the frame before shrinking it, so a
    box near (but not flush against) the border keeps its square aspect and
    minimum size; only a box already pinned against two edges, or a frame
    smaller than the requested box, forces an actual shrink.
    """
    ys, xs = np.nonzero(mask)
    x0f, x1f = float(xs.min()), float(xs.max() + 1)
    y0f, y1f = float(ys.min()), float(ys.max() + 1)

    cx, cy = (x0f + x1f) / 2.0, (y0f + y1f) / 2.0
    bw = (x1f - x0f) * (1.0 + 2.0 * padding)
    bh = (y1f - y0f) * (1.0 + 2.0 * padding)

    if square:
        bw = bh = max(bw, bh)
    bw = max(bw, float(min_size))
    bh = max(bh, float(min_size))

    x0, x1 = cx - bw / 2.0, cx + bw / 2.0
    y0, y1 = cy - bh / 2.0, cy + bh / 2.0

    if x0 < 0:
        x1 -= x0
        x0 = 0.0
    if y0 < 0:
        y1 -= y0
        y0 = 0.0
    if x1 > w:
        x0 -= x1 - w
        x1 = float(w)
    if y1 > h:
        y0 -= y1 - h
        y1 = float(h)

    x0i, y0i = int(np.floor(max(x0, 0.0))), int(np.floor(max(y0, 0.0)))
    x1i, y1i = int(np.ceil(min(x1, w))), int(np.ceil(min(y1, h)))
    return x0i, y0i, x1i, y1i


def select_best_frame(masklets_or_masks: Any, prefer: str = "largest") -> int:
    """Pick the clip frame whose mask makes the best single view of the object.

    Rationale: every downstream stage (crop, image-to-3D generation, alignment)
    only ever sees *this one frame*. Unlike a multi-view reconstruction pipeline
    there is no later stage that can average away a bad choice here, so it is
    worth choosing deliberately rather than defaulting to frame 0 or the frame a
    prompt happened to be applied on.

    Args:
        masklets_or_masks: A :class:`~fpgm.types.Masklet` (uses ``.frames``), a
            ``dict[int, mask]``, or a sequence of masks indexed by frame.
        prefer: ``"largest"`` (mask area is the primary signal, with a frame
            touching the image border penalised rather than excluded -- on
            some clips every frame touches the border slightly) or
            ``"least_truncated"`` (minimise border-touching pixels first, area
            only as a tiebreaker -- use when partial views are worse than a
            smaller full view for the downstream generator).

    Returns:
        The chosen frame index.

    Raises:
        ValueError: If ``prefer`` is unknown or every mask is empty.
    """
    if prefer not in _SELECT_METHODS:
        raise ValueError(
            f"select_best_frame: unknown prefer {prefer!r}, expected one of {_SELECT_METHODS}"
        )

    frames = _coerce_frame_masks(masklets_or_masks)
    frames = {idx: m for idx, m in frames.items() if m is not None and m.any()}
    if not frames:
        raise ValueError("select_best_frame: no non-empty masks to choose from")

    def border_px(m: np.ndarray) -> int:
        return int(m[0, :].sum() + m[-1, :].sum() + m[:, 0].sum() + m[:, -1].sum())

    def area(m: np.ndarray) -> int:
        return int(m.sum())

    if prefer == "largest":
        key = lambda idx: (border_px(frames[idx]) == 0, area(frames[idx]))  # noqa: E731
    else:
        key = lambda idx: (  # noqa: E731
            border_px(frames[idx]) == 0,
            -border_px(frames[idx]),
            area(frames[idx]),
        )
    return max(frames, key=key)


def _coerce_frame_masks(obj: Any) -> dict[int, np.ndarray]:
    if hasattr(obj, "frames") and isinstance(obj.frames, dict):
        return obj.frames
    if isinstance(obj, dict):
        return obj
    if isinstance(obj, (list, tuple)):
        return dict(enumerate(obj))
    raise TypeError(
        "select_best_frame: expected a Masklet, dict[int, mask], or sequence of "
        f"masks, got {type(obj).__name__}"
    )


def occlusion_ratio(mask: np.ndarray, robot_mask: np.ndarray) -> float:
    """Fraction of ``mask`` (the object) covered by ``robot_mask``.

    Args:
        mask: ``(H, W)`` boolean object mask.
        robot_mask: ``(H, W)`` boolean robot/gripper mask, same resolution.

    Returns:
        Fraction in ``[0, 1]``; ``0.0`` if ``mask`` is empty.

    Raises:
        ValueError: If the two masks disagree on shape.
    """
    mask = np.asarray(mask, dtype=bool)
    robot_mask = np.asarray(robot_mask, dtype=bool)
    if mask.shape != robot_mask.shape:
        raise ValueError(f"occlusion_ratio: shape mismatch {mask.shape} vs {robot_mask.shape}")
    area = int(mask.sum())
    if area == 0:
        return 0.0
    return float((mask & robot_mask).sum()) / area


def inpaint_occlusion(
    rgba: np.ndarray, occluded: np.ndarray, method: str = "telea", radius: int = 3
) -> tuple[np.ndarray, float]:
    """Fill an occluded region of an RGBA crop with ``cv2.inpaint``.

    This is a placeholder for a proper diffusion inpaint (e.g. LaMa or a
    Stable-Diffusion inpainting model) ahead of the image-to-3D generator.
    ``cv2.inpaint`` is a PDE-based fill: it smears neighbouring texture across
    the hole rather than hallucinating plausible structure, so it is only
    trustworthy for small gaps at an object's edge. A large occluded fraction
    (see :data:`OCCLUSION_TRUST_THRESHOLD`) should be treated as a reason to
    pick a different frame via :func:`select_best_frame`, not as something this
    function can be trusted to have fixed.

    Args:
        rgba: ``(H, W, 4)`` uint8 crop; alpha marks the object.
        occluded: ``(H, W)`` boolean mask of pixels to fill, same resolution.
        method: ``"telea"`` or ``"ns"`` (Navier-Stokes), the two algorithms
            ``cv2.inpaint`` supports.
        radius: Inpainting neighbourhood radius in pixels.

    Returns:
        ``(filled_rgba, inpainted_fraction)`` where ``inpainted_fraction`` is
        the occluded area as a fraction of the *object* area (``alpha > 0``),
        which is what :attr:`~fpgm.objects.types.ObjectCrop.inpainted_fraction`
        expects.

    Raises:
        ValueError: If ``method`` is unknown or shapes disagree.
    """
    if method not in _INPAINT_METHODS:
        raise ValueError(
            f"inpaint_occlusion: unknown method {method!r}, expected one of "
            f"{tuple(_INPAINT_METHODS)}"
        )
    rgba = np.asarray(rgba)
    occluded = np.asarray(occluded, dtype=bool)
    if rgba.shape[:2] != occluded.shape:
        raise ValueError(
            f"inpaint_occlusion: rgba shape {rgba.shape[:2]} != occluded shape {occluded.shape}"
        )

    object_mask = rgba[..., 3] > 0
    object_area = int(object_mask.sum())
    occluded_object = occluded & object_mask
    fraction = float(occluded_object.sum()) / object_area if object_area else 0.0
    if fraction > OCCLUSION_TRUST_THRESHOLD:
        logger.warning(
            "inpaint_occlusion: %.0f%% of the object is occluded; the cv2.inpaint "
            "fill is not trustworthy at this fraction -- prefer a less-occluded "
            "frame if one is available",
            fraction * 100.0,
        )

    inpaint_mask = occluded_object.astype(np.uint8) * 255
    filled_rgb = cv2.inpaint(rgba[..., :3].copy(), inpaint_mask, radius, _INPAINT_METHODS[method])
    out = rgba.copy()
    out[..., :3] = filled_rgb
    return out, fraction


def save_debug(
    crop: ObjectCrop, out_dir: Path, frame_bgr: np.ndarray | None = None
) -> StageArtifacts:
    """Write the crop PNG, a mask overlay, and a stats JSON for ``crop``.

    Args:
        crop: The crop to dump.
        out_dir: Directory to write into (created if missing).
        frame_bgr: The full frame the crop came from. If given, the overlay
            shows the mask contour and crop bbox on the real frame; if omitted,
            the overlay falls back to a flat colour rendering of just the mask
            (still useful, but without scene context).

    Returns:
        A :class:`~fpgm.objects.types.StageArtifacts` pointing at the written
        files.
    """
    out_dir = ensure_dir(Path(out_dir))
    artifacts = StageArtifacts(stage="crop", directory=out_dir)

    crop_path = out_dir / f"crop_{crop.frame_idx:06d}.png"
    cv2.imwrite(str(crop_path), cv2.cvtColor(crop.rgba, cv2.COLOR_RGBA2BGRA))
    artifacts.add("crop_png", crop_path)

    if frame_bgr is not None:
        overlay = overlay_mask(frame_bgr, crop.mask_full, color_for(0))
        x0, y0, x1, y1 = crop.bbox_xyxy
        cv2.rectangle(overlay, (x0, y0), (x1 - 1, y1 - 1), (255, 255, 255), 1, cv2.LINE_AA)
    else:
        logger.warning("save_debug: no frame_bgr given, writing a mask-only overlay")
        overlay = np.zeros((*crop.mask_full.shape, 3), dtype=np.uint8)
        overlay[crop.mask_full] = color_for(0)
    overlay_path = out_dir / f"mask_overlay_{crop.frame_idx:06d}.png"
    cv2.imwrite(str(overlay_path), overlay)
    artifacts.add("mask_overlay", overlay_path)

    stats = {
        "frame_idx": crop.frame_idx,
        "bbox_xyxy": list(crop.bbox_xyxy),
        "area_px": crop.area_px,
        "crop_shape": list(crop.rgba.shape),
        "inpainted_fraction": crop.inpainted_fraction,
        "prompt": crop.prompt,
    }
    stats_path = out_dir / f"crop_stats_{crop.frame_idx:06d}.json"
    stats_path.write_text(json.dumps(stats, indent=2))
    artifacts.add("stats", stats_path)
    artifacts.stats = stats
    return artifacts