src/fpgm/tracking/sampling.py

"""Pure-numpy query-point sampling inside a segmentation mask.

TAPNext++ needs a fixed set of ``(x, y)`` query points seeded on one frame (see
:mod:`fpgm.tracking.tapnext`). This module chooses those points from a boolean
object mask. It has no torch/tapnet/sam3 dependency at all, so it is fully
unit-testable on synthetic masks without a GPU or model weights.
"""

from __future__ import annotations

import logging

import numpy as np
from scipy import ndimage

logger = logging.getLogger(__name__)

_METHODS = ("farthest", "grid", "random")


def mask_centroid(mask: np.ndarray) -> np.ndarray:
    """Return the pixel centroid of a boolean mask.

    Args:
        mask: ``(H, W)`` boolean array.

    Returns:
        ``(2,)`` float32 array ``[u, v]`` (x, y pixel coordinates).

    Raises:
        ValueError: If ``mask`` has no True pixels.
    """
    ys, xs = np.nonzero(mask)
    if xs.size == 0:
        raise ValueError("mask_centroid: mask is empty")
    return np.array([xs.mean(), ys.mean()], dtype=np.float32)


def points_in_mask(mask: np.ndarray, uv: np.ndarray) -> np.ndarray:
    """Test which ``(x, y)`` pixel points fall inside a boolean mask.

    Used by the depth module to restrict its neighbour search to an object's
    mask (see :class:`~fpgm.depth.base.DepthSource`).

    Args:
        mask: ``(H, W)`` boolean array.
        uv: ``(Q, 2)`` array of ``[x, y]`` pixel coordinates. Membership **floors**:
            pixel ``i`` covers ``[i, i + 1)``. Rounding would count points up to
            half a pixel outside the silhouette as inside, which leaks background
            geometry across object boundaries.

    Returns:
        ``(Q,)`` boolean array; False for points outside the mask bounds.
    """
    uv = np.asarray(uv)
    h, w = mask.shape
    xi = np.floor(uv[:, 0]).astype(np.int64)
    yi = np.floor(uv[:, 1]).astype(np.int64)
    in_bounds = (xi >= 0) & (xi < w) & (yi >= 0) & (yi < h)
    out = np.zeros(uv.shape[0], dtype=bool)
    out[in_bounds] = mask[yi[in_bounds], xi[in_bounds]]
    return out


def sample_points_in_mask(
    mask: np.ndarray,
    n_points: int,
    method: str = "farthest",
    boundary_erosion_px: int = 3,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """Sample query points for TAPNext++ inside a segmentation mask.

    Points are drawn from an eroded copy of ``mask`` rather than the raw mask:
    right at the object boundary is exactly where depth interpolation is least
    reliable (see :class:`~fpgm.depth.base.DepthSource`) and where TAPNext++
    most often drifts onto the background, since a one-pixel segmentation error
    is enough to seed a query outside the object entirely.

    Args:
        mask: ``(H, W)`` boolean array.
        n_points: Number of query points requested.
        method: ``"farthest"`` (farthest-point sampling seeded at the mask
            centroid, for even spatial coverage), ``"grid"`` (a regular lattice
            intersected with the mask), or ``"random"`` (uniform random pixels).
        boundary_erosion_px: Pixels to erode off the mask boundary before
            sampling. If erosion empties the mask entirely (e.g. a very thin or
            small object), falls back to the un-eroded mask with a warning.
        rng: Random generator; a fresh :func:`numpy.random.default_rng` is used
            if not given. Only consulted by ``"random"`` and for centroid-tie
            resolution.

    Returns:
        ``(Q, 2)`` float32 array of ``[x, y]`` pixel coordinates, ``Q <=
        n_points``. Fewer than ``n_points`` are returned (with a warning) if
        the mask cannot support that many distinct pixels.

    Raises:
        ValueError: If ``mask`` is empty or ``method``/``n_points`` are invalid.
    """
    if n_points <= 0:
        raise ValueError(f"sample_points_in_mask: n_points must be positive, got {n_points}")
    if method not in _METHODS:
        raise ValueError(
            f"sample_points_in_mask: unknown method {method!r}, expected one of {_METHODS}"
        )

    mask = np.asarray(mask, dtype=bool)
    if not mask.any():
        raise ValueError("sample_points_in_mask: mask is empty")

    if rng is None:
        rng = np.random.default_rng()

    eroded = _erode(mask, boundary_erosion_px)

    if method == "grid":
        points = _grid_sample(eroded, n_points)
    else:
        ys, xs = np.nonzero(eroded)
        candidates = np.stack([xs, ys], axis=1).astype(np.float64)
        n_available = candidates.shape[0]
        k = min(n_points, n_available)
        if k < n_points:
            logger.warning(
                "sample_points_in_mask: requested %d points but the mask only has "
                "%d pixels after erosion; returning %d",
                n_points,
                n_available,
                k,
            )
        if method == "random":
            idx = rng.choice(n_available, size=k, replace=False)
        else:  # farthest
            idx = _farthest_point_sample(candidates, k)
        points = candidates[idx]

    return points.astype(np.float32)


def _erode(mask: np.ndarray, boundary_erosion_px: int) -> np.ndarray:
    """Erode ``mask`` by ``boundary_erosion_px``, falling back if it empties."""
    if boundary_erosion_px <= 0:
        return mask
    eroded = ndimage.binary_erosion(mask, iterations=boundary_erosion_px)
    if not eroded.any():
        logger.warning(
            "sample_points_in_mask: eroding by %d px emptied the mask (area=%d "
            "px); falling back to the un-eroded mask",
            boundary_erosion_px,
            int(mask.sum()),
        )
        return mask
    return eroded


def _farthest_point_sample(candidates: np.ndarray, k: int) -> np.ndarray:
    """Greedy farthest-point sampling, seeded at the point nearest the centroid.

    Args:
        candidates: ``(N, 2)`` float64 pixel coordinates, ``N >= k``.
        k: Number of points to select.

    Returns:
        ``(k,)`` int64 indices into ``candidates``.
    """
    centroid = candidates.mean(axis=0)
    seed_idx = int(np.argmin(np.sum((candidates - centroid) ** 2, axis=1)))

    selected = np.empty(k, dtype=np.int64)
    selected[0] = seed_idx
    min_dist = np.full(candidates.shape[0], np.inf, dtype=np.float64)
    last = candidates[seed_idx]
    for i in range(1, k):
        d = np.sum((candidates - last) ** 2, axis=1)
        min_dist = np.minimum(min_dist, d)
        nxt = int(np.argmax(min_dist))
        selected[i] = nxt
        min_dist[nxt] = -1.0  # never re-select: guarantees k distinct indices
        last = candidates[nxt]
    return selected


def _grid_sample(mask: np.ndarray, n_points: int) -> np.ndarray:
    """Intersect a regular lattice with ``mask``, refining spacing to fit ``n_points``.

    Starts from a spacing estimate based on mask area / n_points, then shrinks
    the lattice step until it yields at least ``n_points`` mask-interior points
    (or hits sub-pixel spacing), and finally trims evenly down to exactly
    ``n_points`` so the result stays close to a genuine lattice rather than a
    dense cluster.
    """
    ys, xs = np.nonzero(mask)
    y0, y1, x0, x1 = int(ys.min()), int(ys.max()), int(xs.min()), int(xs.max())
    area = xs.size

    step = max(1.0, float(np.sqrt(area / max(n_points, 1))))
    pts = np.stack([xs, ys], axis=1).astype(np.float64)  # degenerate fallback
    for _ in range(25):
        xs_grid = np.arange(x0, x1 + 1, step)
        ys_grid = np.arange(y0, y1 + 1, step)
        gx, gy = np.meshgrid(xs_grid, ys_grid)
        candidates = np.stack([gx.ravel(), gy.ravel()], axis=1)
        inside = points_in_mask(mask, candidates)
        found = candidates[inside]
        if found.shape[0] > 0:
            pts = found
        if found.shape[0] >= n_points or step <= 0.5:
            break
        step /= 1.5

    if pts.shape[0] > n_points:
        idx = np.linspace(0, pts.shape[0] - 1, n_points).astype(np.int64)
        pts = pts[idx]
    elif pts.shape[0] < n_points:
        logger.warning(
            "sample_points_in_mask: grid method only found %d lattice points "
            "inside the mask (requested %d); returning %d",
            pts.shape[0],
            n_points,
            pts.shape[0],
        )
    return pts