src/fpgm/datagen/discovery.py

"""S4: discover moving (non-robot) objects with no prior knowledge of what they are.

The premise this whole module exists to satisfy: the task script ("put brick in
drawer shelf and close drawer") names two objects, but a batch run over other
DROID episodes will not come with names, so nothing downstream is allowed to
special-case "brick" or "drawer". The only thing that generalises is *motion*:
seed a query point everywhere in the frame, track it, throw away whatever moves
only because the robot does (S2's render already tells us exactly where that
is) or does not move at all (the camera is static, so residual image-space
motion after removing the robot really is object motion), and let whatever is
left cluster itself into candidate objects. Two or three clusters are expected
on the demo episode (brick, drawer front, and possibly a spurious third); the
number is a *result* of this stage, not an input to it.

**Design choices made here, and why the alternatives were rejected:**

* **Seed frame 0, not mid-clip.** TAPNext++ is strictly causal and can only be
  seeded once (see :class:`~fpgm.tracking.tapnext.TapNextPointTracker`'s
  docstring): every frame before the seed comes back ``visible=False`` and is
  never back-filled. Seeding mid-episode would lose every earlier frame
  outright -- exactly the frames that show the brick sitting untouched on the
  table and the drawer hanging open, which is precisely the "what does the
  scene look like before anything happens" reference this stage needs. Seeding
  at 0 also avoids the two-rollout complexity the task brief flags as the
  alternative (seed early, seed late, merge): one rollout, full coverage,
  nothing to reconcile. The one risk -- frame 0 being occluded -- was checked,
  not assumed: the S2 robot mask at frame 0 covers 34,686 / 921,600 px (3.8%)
  of the frame, concentrated on the arm base far from the table, and the task
  script's own choreography (brick on the table, drawer already open) means
  both target objects are un-occluded at the start by construction. A frame
  later in the episode risks exactly the opposite: the brick disappearing
  inside the drawer.
* **Net displacement, not path length, for the static-track gate
  (:func:`filter_static_tracks`).** A point whose track wanders under tracking
  jitter and returns close to where it started has zero true object motion,
  even though its path length is nonzero; using path length would keep every
  jittery background point that TAPNext++ merely tracks noisily. Net
  displacement (first-visible position to farthest-visible position) is zero
  for that case and large for genuine sustained motion, which is the
  distinction that actually matters here.
* **Scene-flow points, not S2's robot depth or the per-clip ``initial_depth``
  anchor, as the 3D lift source (:func:`lift_query_points`).** S2's
  ``depth_mm.npy`` is zero everywhere the robot render does not cover -- it
  carries no information at all about the brick or the drawer, which is
  exactly what this stage needs depth for. ``initial_depth`` is dense but
  fixed to each clip's own local frame 0, which is not guaranteed to be this
  stage's seed frame in the episode's video-frame axis. ``scene_flows``
  positions are already real, triangulated, world-frame 3D points (verified
  aligned for this episode -- see the plan doc) available at any annotated
  frame of any covering clip, so a query pixel is matched to its *nearest*
  projected scene-flow point (within :data:`_SCENE_FLOW_MATCH_RADIUS_PX`) and
  that point's own world position is used directly -- no depth-plus-unproject
  round trip, and no dependency on S2 at all (this module never imports
  ``fpgm.datagen.robot_buffers``; the robot mask is taken as a plain injected
  array, the same seam :mod:`fpgm.datagen.dense_depth` uses for S3/S2).
* **DBSCAN over ``[x, y, z, dx_m, dy_m]``, not 2D pixel motion alone.** Pixel
  motion for a fixed real-world displacement is inversely proportional to
  depth (perspective), so two objects moving the same true distance at
  different depths from the camera would otherwise cluster apart on
  displacement magnitude alone for the wrong reason. Converting the net pixel
  displacement to an approximate metric vector via the pinhole relation
  ``d_m = d_px * z / f`` (see :func:`build_cluster_features`) puts position and
  displacement in the same physical units, so a single ``eps`` (in metres,
  already the unit ``DatagenConfig.dbscan_eps_m`` is documented in) applies
  meaningfully to both.
"""

from __future__ import annotations

from dataclasses import dataclass, replace
from pathlib import Path

import cv2
import numpy as np

from fpgm.config import ConventionConfig, DatagenConfig, TrackingConfig
from fpgm.data.pointworld import FlowsReader
from fpgm.datagen.cache import StageCache
from fpgm.datagen.frame_index import EpisodeFrameIndex
from fpgm.geometry.camera import Camera
from fpgm.geometry.convention import detect_scene_flow_convention
from fpgm.types import DataError, FrameConvention, Track2D
from fpgm.utils.io import ensure_dir
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

#: StageCache stage name this module writes under.
STAGE_NAME = "discovery"

# TODO(config): these three thresholds belong on DatagenConfig (alongside its
# existing tapnext_grid/dbscan_eps_m/dbscan_min_samples fields), but config.py
# is sibling-owned and out of scope here. Kept as module-level constants with
# their rationale recorded in the module docstring; a caller who wants a
# different value passes it explicitly to the functions below rather than
# editing this file.
_ROBOT_LIFE_FRACTION_GATE = 0.5  # drop a track if >= half its visible life is inside the robot mask
_STATIC_DISPLACEMENT_GATE_PX = 6.0  # net displacement below this is "the camera's own noise floor"
_SCENE_FLOW_MATCH_RADIUS_PX = 15.0  # nearest-scene-flow-point match radius for the 3D lift


# --------------------------------------------------------------------------- #
# Query grid
# --------------------------------------------------------------------------- #


def sample_grid_points(
    width: int,
    height: int,
    n_cols: int,
    n_rows: int,
    jitter_frac: float = 0.0,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """A full-frame, cell-centred query grid for seeding TAPNext++.

    Cell-*centred* rather than edge-aligned: a lattice that starts at pixel
    ``(0, 0)`` puts a quarter of its points within a few pixels of the frame
    border, exactly where a tracker is most likely to lose a point to the edge
    on the very first frame and where TAPNext++'s own receptive field is
    truncated. Centring each cell avoids seeding on that unstable margin for
    free, at zero cost to coverage.

    ``jitter_frac``/``rng`` exist only so a caller who wants to break up a
    perfectly periodic lattice (e.g. to avoid every query landing on the same
    phase of a repeating background texture) can do so reproducibly -- with
    ``jitter_frac=0`` (the default) the result has no randomness in it at all.

    Args:
        width: Frame width, pixels.
        height: Frame height, pixels.
        n_cols: Number of grid columns.
        n_rows: Number of grid rows.
        jitter_frac: Fraction of one cell's size to jitter each point by,
            uniformly in ``[-0.5, 0.5] * jitter_frac * cell_size``. ``0``
            (default) disables jitter entirely.
        rng: Random generator for jitter; a fresh
            :func:`numpy.random.default_rng` is used if omitted. Ignored when
            ``jitter_frac == 0``.

    Returns:
        ``(n_cols * n_rows, 2)`` float32 array of ``[x, y]`` pixel coordinates.

    Raises:
        ValueError: If ``width``, ``height``, ``n_cols``, or ``n_rows`` is not
            positive.
    """
    if width <= 0 or height <= 0:
        raise ValueError(
            f"sample_grid_points: width/height must be positive, got {width}x{height}"
        )
    if n_cols <= 0 or n_rows <= 0:
        raise ValueError(
            f"sample_grid_points: n_cols/n_rows must be positive, got {n_cols}x{n_rows}"
        )

    cell_w = width / n_cols
    cell_h = height / n_rows
    xs = (np.arange(n_cols, dtype=np.float64) + 0.5) * cell_w
    ys = (np.arange(n_rows, dtype=np.float64) + 0.5) * cell_h
    gx, gy = np.meshgrid(xs, ys)
    pts = np.stack([gx.ravel(), gy.ravel()], axis=1)

    if jitter_frac > 0:
        rng = rng or np.random.default_rng()
        jitter = (rng.random(pts.shape) - 0.5) * jitter_frac * np.array([cell_w, cell_h])
        pts = pts + jitter
        pts[:, 0] = np.clip(pts[:, 0], 0.0, width - 1.0)
        pts[:, 1] = np.clip(pts[:, 1], 0.0, height - 1.0)

    return pts.astype(np.float32)


# --------------------------------------------------------------------------- #
# Track2D (de)serialization -- shared with object_masks.py
# --------------------------------------------------------------------------- #


def write_track2d(path: str | Path, track: Track2D) -> Path:
    """Write a :class:`~fpgm.types.Track2D` to a compressed ``.npz``."""
    path = Path(path)
    ensure_dir(path.parent)
    np.savez_compressed(
        path,
        point_id=track.point_id,
        frames=track.frames,
        uv=track.uv,
        visible=track.visible,
        resolution=np.array(track.resolution, dtype=np.int64),
    )
    return path


def read_track2d(path: str | Path) -> Track2D:
    """Inverse of :func:`write_track2d`."""
    with np.load(path) as data:
        return Track2D(
            point_id=data["point_id"],
            frames=data["frames"],
            uv=data["uv"],
            visible=data["visible"],
            resolution=tuple(int(x) for x in data["resolution"]),
        )


def _subset_track(track: Track2D, keep: np.ndarray) -> Track2D:
    """Restrict a :class:`Track2D` to a subset of its query-point (Q) axis."""
    keep = np.asarray(keep, dtype=bool)
    return Track2D(
        point_id=track.point_id[keep],
        frames=track.frames,
        uv=track.uv[:, keep, :],
        visible=track.visible[:, keep],
        resolution=track.resolution,
    )


# --------------------------------------------------------------------------- #
# Robot subtraction
# --------------------------------------------------------------------------- #


def compute_robot_overlap(
    track: Track2D, robot_seg: np.ndarray, query_frame_idx: int
) -> tuple[np.ndarray, np.ndarray]:
    """Per-track robot-mask overlap: at the seed frame, and across its visible life.

    Args:
        track: Tracks to test, ``uv``/``visible`` shaped ``(T, Q, 2)``/``(T, Q)``.
        robot_seg: ``(T', H, W)`` integer array, nonzero = robot (S2's
            ``robot_buffers/seg.npy``, mmap-friendly -- indexed one frame at a
            time so a caller passing an ``np.load(..., mmap_mode="r")`` array
            never materialises more than one frame here).
        query_frame_idx: Absolute video frame the tracks were seeded on.

    Returns:
        ``(at_query_frame, life_fraction)``:

        * ``at_query_frame``: ``(Q,)`` bool, whether the query point itself
          landed inside the robot mask on the seed frame -- this is the
          decisive case (a point seeded on the arm can never be anything but
          "the arm moved"), checked separately from the life-fraction so a
          caller can gate on it alone if desired.
        * ``life_fraction``: ``(Q,)`` float32 in ``[0, 1]``, fraction of each
          track's *visible* frames spent inside the robot mask. ``0`` for a
          track with zero visible frames (nothing to divide by, and nothing to
          accuse of being the robot either).
    """
    n_frames_track, n_points, _ = track.uv.shape
    hits = np.zeros(n_points, dtype=np.int64)
    totals = np.zeros(n_points, dtype=np.int64)
    at_query = np.zeros(n_points, dtype=bool)

    for t in range(n_frames_track):
        frame = int(track.frames[t])
        vis = track.visible[t]
        if not vis.any():
            continue
        if frame < 0 or frame >= robot_seg.shape[0]:
            continue
        seg_frame = np.asarray(robot_seg[frame]) > 0
        h, w = seg_frame.shape
        uv = track.uv[t]
        xi = np.floor(uv[:, 0]).astype(np.int64)
        yi = np.floor(uv[:, 1]).astype(np.int64)
        in_bounds = vis & (xi >= 0) & (xi < w) & (yi >= 0) & (yi < h)
        is_robot = np.zeros(n_points, dtype=bool)
        is_robot[in_bounds] = seg_frame[yi[in_bounds], xi[in_bounds]]

        totals[vis] += 1
        hits[is_robot] += 1
        if frame == query_frame_idx:
            at_query = is_robot

    life_fraction = np.divide(
        hits, totals, out=np.zeros(n_points, dtype=np.float64), where=totals > 0
    ).astype(np.float32)
    return at_query, life_fraction


def filter_robot_tracks(
    track: Track2D,
    robot_seg: np.ndarray,
    query_frame_idx: int,
    life_fraction_gate: float = _ROBOT_LIFE_FRACTION_GATE,
) -> np.ndarray:
    """Which of ``track``'s query points are NOT the robot.

    See :func:`compute_robot_overlap` for the two signals combined here.
    Either one alone is not enough: a point seeded just off the arm can drift
    onto it later without ever being "seeded on the robot" (life-fraction
    catches that), while a point seeded exactly on the arm might legitimately
    spend under half its life there if the arm sweeps away and something else
    passes under the query pixel afterward (the query-frame check catches
    that, independent of the life-fraction threshold).

    Returns:
        ``(Q,)`` bool, ``True`` = keep (not the robot).
    """
    at_query, life_fraction = compute_robot_overlap(track, robot_seg, query_frame_idx)
    return (~at_query) & (life_fraction < life_fraction_gate)


# --------------------------------------------------------------------------- #
# Static-track rejection
# --------------------------------------------------------------------------- #


def compute_displacement(track: Track2D) -> tuple[np.ndarray, np.ndarray]:
    """Net 2D displacement of each track, from its first to its farthest visible position.

    See the module docstring for why net displacement (not path length) is the
    right static/dynamic discriminator here.

    Returns:
        ``(net_vector_px, magnitude_px)``:

        * ``net_vector_px``: ``(Q, 2)`` float32, ``uv[farthest] - uv[first_visible]``.
          ``[0, 0]`` for a track with fewer than 2 visible frames.
        * ``magnitude_px``: ``(Q,)`` float32, ``||net_vector_px||``.
    """
    n_frames, n_points, _ = track.uv.shape
    net = np.zeros((n_points, 2), dtype=np.float32)
    for q in range(n_points):
        vis_idx = np.flatnonzero(track.visible[:, q])
        if vis_idx.size < 2:
            continue
        pts = track.uv[vis_idx, q]
        start = pts[0]
        dists = np.linalg.norm(pts - start[None, :], axis=1)
        far = int(np.argmax(dists))
        net[q] = pts[far] - start
    magnitude = np.linalg.norm(net, axis=1).astype(np.float32)
    return net, magnitude


def filter_static_tracks(
    track: Track2D, displacement_gate_px: float = _STATIC_DISPLACEMENT_GATE_PX
) -> np.ndarray:
    """Which of ``track``'s query points moved enough to be considered dynamic.

    Returns:
        ``(Q,)`` bool, ``True`` = keep (moved at least ``displacement_gate_px``).
    """
    _, magnitude = compute_displacement(track)
    return magnitude >= displacement_gate_px


# --------------------------------------------------------------------------- #
# 3D lift via nearest scene-flow point
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class _SceneFlowAnchor:
    """One frame's scene-flow point cloud, projected and ready to match query pixels against."""

    world_xyz: np.ndarray  # (N, 3) float64, world frame, metres
    uv: np.ndarray  # (N, 2) float64, pixels at the anchor's own (video) resolution
    z_cam: np.ndarray  # (N,) float64, camera-frame depth, metres
    camera: Camera  # rescaled to the anchor's resolution


def _build_scene_flow_anchor(
    frame_index: EpisodeFrameIndex,
    flows_reader: FlowsReader,
    camera_serial: str,
    query_frame_idx: int,
    video_width: int,
    video_height: int,
    native_resolution: tuple[int, int],
    convention_cfg: ConventionConfig,
) -> _SceneFlowAnchor:
    """Build the scene-flow point cloud for ``query_frame_idx``, at video resolution.

    Tries every clip covering ``query_frame_idx`` in order (there can be
    several -- clips overlap by design) and uses the first one whose local
    frame actually has visible, depth-valid scene-flow points; a clip that
    covers the frame in principle but has nothing usable there (e.g. every
    point occluded) is skipped rather than failing the whole lift.

    Raises:
        DataError: If no covering clip has usable scene-flow data at
            ``query_frame_idx`` -- the 3D lift genuinely cannot proceed, and
            silently returning an empty anchor would make every downstream
            track look statically un-liftable instead of surfacing the real
            cause.
    """
    native_w, native_h = native_resolution
    covering = frame_index.clips_covering(query_frame_idx)
    if not covering:
        raise DataError(
            f"no clip covers video frame {query_frame_idx} for camera {camera_serial!r}; "
            "cannot build a scene-flow anchor to lift discovery tracks to 3D"
        )

    for clip_ref in covering:
        t_local = frame_index.video_to_clip_local(clip_ref.key, query_frame_idx)
        if t_local is None:
            continue
        sfc = flows_reader.read_clip(clip_ref.key, camera_serial)
        if sfc.initial_rgb is None:
            logger.warning(
                "discovery: clip %s has no initial_rgb, cannot detect scene_flows "
                "convention; skipping as an anchor source",
                clip_ref.key,
            )
            continue
        camera_native = Camera.from_pointworld(sfc.intrinsic, sfc.extrinsic, native_w, native_h)
        detection = detect_scene_flow_convention(
            camera_native, sfc.scene_flows[0], sfc.scene_colors[0], sfc.initial_rgb, convention_cfg
        )
        frame_is_world = detection.convention is FrameConvention.WORLD
        camera = camera_native.rescaled(video_width, video_height)

        pts = sfc.scene_flows[t_local].astype(np.float64)
        visible = np.asarray(sfc.scene_visibility[t_local])
        depth_valid = np.asarray(sfc.scene_depth_valid[t_local])
        pts = pts[visible & depth_valid]
        if pts.shape[0] == 0:
            continue

        if frame_is_world:
            uv, z = camera.project(pts)
            world_xyz = pts
        else:
            uv, z = camera.project_cam(pts)
            world_xyz = camera.cam_to_world(pts)

        in_front = z > 0
        if not in_front.any():
            continue
        return _SceneFlowAnchor(
            world_xyz=world_xyz[in_front], uv=uv[in_front], z_cam=z[in_front], camera=camera
        )

    raise DataError(
        f"every clip covering video frame {query_frame_idx} (camera {camera_serial!r}) had "
        "no usable (visible AND depth-valid) scene-flow points -- cannot lift discovery "
        "tracks to 3D at this frame"
    )


def lift_query_points(
    anchor: _SceneFlowAnchor, query_uv: np.ndarray, max_radius_px: float
) -> tuple[np.ndarray, np.ndarray]:
    """Nearest-scene-flow-point 3D lift for a set of 2D query pixels.

    Args:
        anchor: A :class:`_SceneFlowAnchor` built at the query points' frame
            and resolution.
        query_uv: ``(Q, 2)`` float pixel coordinates.
        max_radius_px: Points farther than this from every scene-flow
            projection are left unlifted (``NaN``) rather than matched to a
            distant, unrelated point -- a query pixel with no nearby
            scene-flow support (e.g. it landed on a texture-poor patch that
            PointWorld never tracked) should say so honestly, not silently
            borrow whatever the nearest annotated point happens to be.

    Returns:
        ``(world_xyz, z_cam)``: ``(Q, 3)`` float64 (``NaN`` rows where
        unmatched) and ``(Q,)`` float64 (``NaN`` where unmatched).
    """
    from scipy.spatial import cKDTree

    query_uv = np.asarray(query_uv, dtype=np.float64)
    n = query_uv.shape[0]
    world = np.full((n, 3), np.nan, dtype=np.float64)
    z = np.full(n, np.nan, dtype=np.float64)
    if anchor.uv.shape[0] == 0:
        return world, z

    tree = cKDTree(anchor.uv)
    dist, idx = tree.query(query_uv, k=1)
    matched = dist <= max_radius_px
    world[matched] = anchor.world_xyz[idx[matched]]
    z[matched] = anchor.z_cam[idx[matched]]
    return world, z


# --------------------------------------------------------------------------- #
# Clustering
# --------------------------------------------------------------------------- #


def build_cluster_features(
    world_xyz: np.ndarray, net_vector_px: np.ndarray, z_cam: np.ndarray, fx: float, fy: float
) -> np.ndarray:
    """Combine 3D position + metric-approximate displacement into one DBSCAN feature.

    See the module docstring for why the pixel displacement is converted to
    metres via the pinhole relation ``d_m = d_px * z / f`` rather than used
    directly.

    Args:
        world_xyz: ``(N, 3)`` float64, world-frame metres.
        net_vector_px: ``(N, 2)`` float, net pixel displacement (see
            :func:`compute_displacement`).
        z_cam: ``(N,)`` float, camera-frame depth at each point, metres.
        fx: Camera focal length (x), pixels, at the resolution ``net_vector_px``
            was measured at.
        fy: Camera focal length (y), pixels, same resolution.

    Returns:
        ``(N, 5)`` float64: ``[x, y, z, dx_m, dy_m]``.
    """
    world_xyz = np.asarray(world_xyz, dtype=np.float64)
    net_vector_px = np.asarray(net_vector_px, dtype=np.float64)
    z_cam = np.asarray(z_cam, dtype=np.float64)
    dx_m = net_vector_px[:, 0] * z_cam / fx
    dy_m = net_vector_px[:, 1] * z_cam / fy
    return np.concatenate([world_xyz, np.stack([dx_m, dy_m], axis=1)], axis=1)


def cluster_tracks(features: np.ndarray, eps_m: float, min_samples: int) -> np.ndarray:
    """DBSCAN over pre-built ``(N, D)`` feature vectors.

    DBSCAN (not k-means or a fixed cluster count) is the right tool here
    specifically because the number of dynamic objects in an arbitrary DROID
    episode is *unknown* -- that is the entire point of this stage. DBSCAN's
    density threshold (``min_samples`` real neighbours within ``eps``) also
    gives the "does not hallucinate a cluster" property this stage's tests
    require for free: a scattering of noisy, uncorrelated tracks with no real
    spatial/motion agreement produces only noise points (label ``-1``), never
    a spurious cluster, because no point in pure noise has enough neighbours
    within ``eps`` to seed one.

    Args:
        features: ``(N, D)`` float array, already in physically comparable
            units (see :func:`build_cluster_features`) so a single ``eps``
            applies meaningfully across every feature dimension.
        eps_m: DBSCAN neighbourhood radius, in the same units as ``features``.
        min_samples: DBSCAN core-point neighbour threshold.

    Returns:
        ``(N,)`` int32 cluster labels; ``-1`` = noise (not assigned to any
        cluster), following scikit-learn's own convention.
    """
    from sklearn.cluster import DBSCAN

    features = np.asarray(features, dtype=np.float64)
    if features.shape[0] == 0:
        return np.zeros((0,), dtype=np.int32)
    labels = DBSCAN(eps=eps_m, min_samples=min_samples).fit(features).labels_
    return labels.astype(np.int32)


# --------------------------------------------------------------------------- #
# Result assembly
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class DiscoveredObject:
    """One DBSCAN cluster of surviving tracks: a candidate dynamic object.

    Carries only what later stages need to prompt a segmenter at this
    object's location (:mod:`fpgm.datagen.object_masks`) and what a human
    needs to sanity-check the cluster -- not the object's identity, which this
    stage deliberately never determines.
    """

    cluster_id: int
    query_uv: np.ndarray  # (n, 2) float32, pixels at the seed frame -- this cluster's member tracks
    #: (T, 2) float32, mean pixel position of visible members per frame; NaN
    #: where no member was visible that frame.
    centroid_uv_per_frame: np.ndarray
    n_tracks: int
    total_displacement_px: float  # mean net-displacement magnitude across members
    mean_3d_position: np.ndarray  # (3,) float64, world metres, mean over lifted members
    #: (2,) float64, mean metric [dx_m, dy_m] net displacement across members (same
    #: pinhole conversion as build_cluster_features) -- direction + magnitude in
    #: physical units, which is what merge_consistent_clusters compares clusters on.
    #: Defaulted (not required) so existing callers/tests that construct this
    #: dataclass directly without motion data keep working unchanged.
    mean_displacement_m: np.ndarray = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        query_uv = np.asarray(self.query_uv)
        if query_uv.ndim != 2 or query_uv.shape[1] != 2:
            raise DataError(
                f"DiscoveredObject {self.cluster_id}: query_uv must be (n, 2), got "
                f"{query_uv.shape}"
            )
        centroid = np.asarray(self.centroid_uv_per_frame)
        if centroid.ndim != 2 or centroid.shape[1] != 2:
            raise DataError(
                f"DiscoveredObject {self.cluster_id}: centroid_uv_per_frame must be (T, 2), "
                f"got {centroid.shape}"
            )
        mean_3d = np.asarray(self.mean_3d_position)
        if mean_3d.shape != (3,):
            raise DataError(
                f"DiscoveredObject {self.cluster_id}: mean_3d_position must be (3,), got "
                f"{mean_3d.shape}"
            )
        if self.mean_displacement_m is None:
            object.__setattr__(self, "mean_displacement_m", np.zeros(2, dtype=np.float64))
        mean_disp = np.asarray(self.mean_displacement_m)
        if mean_disp.shape != (2,):
            raise DataError(
                f"DiscoveredObject {self.cluster_id}: mean_displacement_m must be (2,), got "
                f"{mean_disp.shape}"
            )


@dataclass(frozen=True)
class DiscoveryStats:
    """Per-stage funnel counts, for the human report: how many points survived each filter."""

    n_query_points: int
    n_after_robot_filter: int
    n_after_static_filter: int
    n_with_3d_lift: int
    n_clusters: int
    n_noise: int


@dataclass(frozen=True)
class DiscoveryResult:
    """S4's output: the surviving tracks, their cluster assignment, and the clusters themselves."""

    objects: tuple[DiscoveredObject, ...]
    #: post robot/static/lift filtering -- every point here fed the clustering.
    track: Track2D
    cluster_labels: np.ndarray  # (Q,) int32, aligned with `track`'s point axis; -1 = noise
    stats: DiscoveryStats
    seed_frame_idx: int

    def __post_init__(self) -> None:
        q = self.track.uv.shape[1]
        labels = np.asarray(self.cluster_labels)
        if labels.shape != (q,):
            raise DataError(
                f"DiscoveryResult: cluster_labels shape {labels.shape} != (Q,) = ({q},) "
                "from `track`'s own point axis -- labels must stay aligned with the "
                "surviving tracks they were computed from"
            )


def _recombine_uv_stats(
    track: Track2D, labels: np.ndarray, cluster_id: int
) -> tuple[np.ndarray, np.ndarray]:
    """``(query_uv, centroid_uv_per_frame)`` for ``cluster_id`` from ``track``/``labels`` alone.

    Split out of :func:`_assemble_objects` so :func:`merge_consistent_clusters` can
    recompute these two fields *exactly* for a merged cluster (a plain relabel
    of the raw per-point ``labels`` array, then a call here) without needing
    the per-point ``world_xyz``/displacement arrays that produced the
    *other* :class:`DiscoveredObject` fields -- those are not carried in
    :class:`DiscoveryResult` (too large to keep around after this stage), so
    the merge path recombines them via the exact weighted-mean identity
    instead (see that function's docstring).
    """
    seed_idx_candidates = np.flatnonzero(track.frames == track.frames[0])
    # query_uv is taken at this track's own first frame, which is always
    # the seed frame for a whole-video TAPNext rollout (frame_indices are
    # contiguous from 0) -- seed_idx_candidates[0] rather than assuming 0
    # so a caller who handed in a track windowed differently still works.
    q0 = int(seed_idx_candidates[0]) if seed_idx_candidates.size else 0
    member = labels == cluster_id
    query_uv = track.uv[q0][member].astype(np.float32)

    n_frames = track.uv.shape[0]
    vis = track.visible[:, member]
    uv = track.uv[:, member, :]
    counts = vis.sum(axis=1)
    sums = np.where(vis[..., None], uv, 0.0).sum(axis=1)
    centroid = np.full((n_frames, 2), np.nan, dtype=np.float32)
    has_any = counts > 0
    centroid[has_any] = (sums[has_any] / counts[has_any, None]).astype(np.float32)
    return query_uv, centroid


def _assemble_objects(
    track: Track2D,
    labels: np.ndarray,
    world_xyz: np.ndarray,
    displacement_px: np.ndarray,
    displacement_m: np.ndarray,
) -> tuple[DiscoveredObject, ...]:
    objects: list[DiscoveredObject] = []
    for cluster_id in sorted(set(int(x) for x in labels.tolist()) - {-1}):
        member = labels == cluster_id
        query_uv, centroid = _recombine_uv_stats(track, labels, cluster_id)
        mean_3d = np.nanmean(world_xyz[member], axis=0)
        mean_disp_m = np.mean(displacement_m[member], axis=0)
        objects.append(
            DiscoveredObject(
                cluster_id=cluster_id,
                query_uv=query_uv,
                centroid_uv_per_frame=centroid,
                n_tracks=int(member.sum()),
                total_displacement_px=float(np.mean(displacement_px[member])),
                mean_3d_position=mean_3d.astype(np.float64),
                mean_displacement_m=mean_disp_m.astype(np.float64),
            )
        )
    return tuple(objects)


def discover_objects(
    track: Track2D,
    robot_seg: np.ndarray,
    query_frame_idx: int,
    world_xyz: np.ndarray,
    z_cam: np.ndarray,
    fx: float,
    fy: float,
    cfg: DatagenConfig,
    robot_life_fraction_gate: float = _ROBOT_LIFE_FRACTION_GATE,
    static_displacement_gate_px: float = _STATIC_DISPLACEMENT_GATE_PX,
) -> DiscoveryResult:
    """Pure function: full-frame grid tracks -> filtered, clustered, dynamic objects.

    Deliberately takes ``world_xyz``/``z_cam`` (already lifted, e.g. via
    :func:`lift_query_points`) rather than the h5/camera machinery needed to
    build them, so this -- the actual discovery *logic* -- is unit-testable on
    synthetic tracks with no GPU, no network, and no PointWorld data. See
    :class:`DynamicObjectStage` for the orchestration that supplies real
    inputs.

    Args:
        track: Full-video, whole-grid tracks from TAPNext++ (``Q`` = the
            original grid size).
        robot_seg: ``(T, H, W)`` integer array, nonzero = robot.
        query_frame_idx: Absolute video frame the tracks were seeded on.
        world_xyz: ``(Q, 3)`` float64, world-frame metres, aligned with
            ``track``'s point axis; ``NaN`` rows are excluded from clustering.
        z_cam: ``(Q,)`` float64, camera-frame depth at each query point.
        fx: Camera focal length (x), pixels.
        fy: Camera focal length (y), pixels.
        cfg: Supplies ``dbscan_eps_m``/``dbscan_min_samples``.
        robot_life_fraction_gate: See :func:`filter_robot_tracks`.
        static_displacement_gate_px: See :func:`filter_static_tracks`.

    Returns:
        A :class:`DiscoveryResult`.
    """
    n_query = track.uv.shape[1]

    keep_robot = filter_robot_tracks(track, robot_seg, query_frame_idx, robot_life_fraction_gate)
    track_r = _subset_track(track, keep_robot)
    world_r = np.asarray(world_xyz)[keep_robot]
    z_r = np.asarray(z_cam)[keep_robot]

    keep_static = filter_static_tracks(track_r, static_displacement_gate_px)
    track_s = _subset_track(track_r, keep_static)
    world_s = world_r[keep_static]
    z_s = z_r[keep_static]
    net_s, mag_s = compute_displacement(track_s)

    has_lift = np.all(np.isfinite(world_s), axis=1) & np.isfinite(z_s)
    track_l = _subset_track(track_s, has_lift)
    world_l = world_s[has_lift]
    z_l = z_s[has_lift]
    net_l = net_s[has_lift]
    mag_l = mag_s[has_lift]

    if track_l.uv.shape[1] > 0:
        features = build_cluster_features(world_l, net_l, z_l, fx, fy)
        labels = cluster_tracks(features, cfg.dbscan_eps_m, cfg.dbscan_min_samples)
        disp_m_l = features[:, 3:5]  # [dx_m, dy_m], same columns build_cluster_features built
    else:
        labels = np.zeros((0,), dtype=np.int32)
        disp_m_l = np.zeros((0, 2), dtype=np.float64)

    objects = _assemble_objects(track_l, labels, world_l, mag_l, disp_m_l)
    n_clusters = len({int(x) for x in labels.tolist()} - {-1})
    n_noise = int(np.count_nonzero(labels == -1))

    stats = DiscoveryStats(
        n_query_points=n_query,
        n_after_robot_filter=int(keep_robot.sum()),
        n_after_static_filter=int(keep_static.sum()),
        n_with_3d_lift=int(has_lift.sum()),
        n_clusters=n_clusters,
        n_noise=n_noise,
    )
    return DiscoveryResult(
        objects=objects,
        track=track_l,
        cluster_labels=labels,
        stats=stats,
        seed_frame_idx=query_frame_idx,
    )


# --------------------------------------------------------------------------- #
# Cluster merging (over-segmentation cleanup)
# --------------------------------------------------------------------------- #

#: Measured on the demo episode's retuned (80x45 grid, eps=0.04, min_samples=3)
#: run: the drawer's two dominant raw DBSCAN clusters (n=82, n=165) sit
#: **0.192 m** apart in mean 3D position -- a wide drawer front genuinely
#: spans that much, since DBSCAN's tight eps (needed to resolve the small
#: brick) fragments it wherever member density dips below eps along its own
#: width. 0.25 m covers that measured gap with margin, while the brick sits
#: ~0.20-0.27 m from either drawer fragment -- close enough that position
#: alone cannot be trusted to separate them at this radius. That separation
#: is instead the direction gate's job (see
#: :data:`_MERGE_DIRECTION_COS_MIN`): the drawer's two fragments agreed at
#: cos~0.96 (nearly parallel -- a rigid front translates as one piece),
#: while the brick's own displacement pointed in a different enough
#: direction (being picked up, not dragged sideways with the drawer) that a
#: pairwise check against either drawer fragment failed the direction gate
#: even before position was checked. Both gates are load-bearing here, not
#: just position -- see the merge rule in :func:`merge_consistent_clusters`.
_MERGE_POSITION_RADIUS_M = 0.25
#: cos(45 deg) ~ 0.71 -- two patches of one rigid body translating together
#: agree far more tightly than this in practice (the demo drawer's two
#: dominant clusters measured cos ~0.96), so this is a loose gate that
#: mainly rejects two objects that happen to be adjacent but move in
#: unrelated directions (measured ~-0.9 to -0.99 between the brick and
#: either drawer fragment on the demo episode -- moving apart, not together).
_MERGE_DIRECTION_COS_MIN = 0.7
#: A 2.5x speed disagreement is far more than tracking noise on two patches
#: of the same rigid body (which differ only by TAPNext++ jitter) -- the
#: demo drawer's two dominant clusters measured a 2.375x ratio, close to
#: this bound because the smaller, farther-from-the-hinge fragment
#: genuinely swept a shorter arc than the larger one -- but still well
#: under what two independently-moving objects would show if they happened
#: to align in direction by chance.
_MERGE_MAGNITUDE_RATIO_MAX = 2.5


def merge_consistent_clusters(
    result: DiscoveryResult,
    position_radius_m: float = _MERGE_POSITION_RADIUS_M,
    direction_cos_min: float = _MERGE_DIRECTION_COS_MIN,
    magnitude_ratio_max: float = _MERGE_MAGNITUDE_RATIO_MAX,
) -> tuple[DiscoveryResult, tuple[tuple[int, ...], ...]]:
    """Fold DBSCAN's over-segmentation of one rigid body back into one cluster.

    **The problem this exists to fix.** DBSCAN in :func:`discover_objects`
    clusters on ``[x, y, z, dx_m, dy_m]`` with a single, fairly tight
    ``eps_m`` (tuned for the brick's small footprint -- see the module's own
    ``dbscan_eps_m``/``dbscan_min_samples`` docs on :class:`~fpgm.config.
    DatagenConfig`). A large, spatially extended rigid object like the
    drawer front routinely has query points more than ``eps_m`` apart within
    the *same* physical part, so DBSCAN legitimately (by its own density
    rule) splits it into several density-connected patches. Measured on the
    demo episode: the drawer came back as two separate clusters (n=82 and
    n=165 tracks) that both matched the same hand-labelled drawer reference
    mask. Retuning ``eps_m`` upward to swallow this would also swallow the
    ~0.6 m gap to the brick, so the fix has to happen *after* clustering, on
    cluster-level evidence DBSCAN itself does not use: whether the clusters'
    *motion* agrees.

    **The merge rule.** Two clusters merge iff *both*:

    1. their :attr:`DiscoveredObject.mean_3d_position` are within
       ``position_radius_m`` metres of each other ("adjacent in 3D"), and
    2. their :attr:`DiscoveredObject.mean_displacement_m` vectors agree in
       direction (cosine similarity >= ``direction_cos_min``) and magnitude
       (the larger norm is at most ``magnitude_ratio_max`` times the
       smaller) -- "consistent motion".

    Neither test alone is safe: position alone would merge two genuinely
    different objects that happen to sit close together (the brick resting
    near the open drawer before the pick); motion direction alone, with no
    position gate, would merge two unrelated objects that momentarily move
    the same way by coincidence. Both together is exactly "this looks like
    one rigid body, sampled twice" -- which is what over-segmentation is.

    A cluster with zero net displacement (should not occur here, since
    :func:`discover_objects` already dropped static tracks, but a merged
    *group*'s residual member could theoretically land on the boundary) is
    never merged into anything -- a zero vector has no defined direction, so
    the cosine test is skipped for it and it merges into nothing.

    **Union-find, not a fresh DBSCAN pass, over cluster pairs.** The
    pairwise merge relation is not guaranteed transitive: a drawer split
    into three collinear patches (A-B-C) can have A agree with B, B agree
    with C, but A and C -- at opposite ends of a wide drawer front -- sit
    just outside ``position_radius_m`` of each other directly. Requiring
    every pair in a group to pass directly would then merge A+B and B+C but
    leave the drawer as two "clusters" that still share every member point
    with each other's report, which is not a fix. Union-find closes this
    transitively from any chain of pairwise agreements, with only
    O(n_clusters^2) pairwise comparisons -- cheap, since the surviving
    cluster count here is always small (single digits).

    **Exact stat recombination, no re-clustering of raw points.**
    :class:`DiscoveryResult` does not retain each surviving query point's
    lifted ``world_xyz`` (too large to keep around past this stage -- see
    :class:`DiscoveryStats`'s own docstring on what *is* kept). Re-deriving
    a merged group's :attr:`~DiscoveredObject.mean_3d_position`/
    :attr:`~DiscoveredObject.mean_displacement_m`/
    :attr:`~DiscoveredObject.total_displacement_px` therefore uses the
    *arithmetic-mean identity* instead of re-touching raw points: for simple
    per-track means, the union's mean equals the ``n_tracks``-weighted mean
    of the constituent objects' own means -- exact, not an approximation,
    because every one of these fields already is a plain mean over its
    cluster's member tracks. ``query_uv``/``centroid_uv_per_frame`` are
    genuinely position-dependent per frame, so those two *are* recomputed
    from the retained ``result.track``/``result.cluster_labels`` (a cheap
    relabel + :func:`_recombine_uv_stats`, no world-frame data needed).

    Args:
        result: A :class:`DiscoveryResult` from :func:`discover_objects`.
        position_radius_m: See "the merge rule" above.
        direction_cos_min: See "the merge rule" above.
        magnitude_ratio_max: See "the merge rule" above.

    Returns:
        ``(merged_result, merge_groups)``: ``merged_result`` has the same
        ``track``/``seed_frame_idx`` as ``result`` (raw tracks are never
        discarded, only relabelled) but recomputed ``objects``/
        ``cluster_labels``/``stats.n_clusters``; ``result.stats``'s other
        funnel counts (``n_query_points`` etc.) are unaffected by merging
        and are copied through unchanged. ``merge_groups`` is one tuple of
        original ``cluster_id``\\ s per output cluster (singletons included,
        for a complete before/after account), ordered to match the new
        cluster ids (``merge_groups[i]`` produced output cluster ``i``).
    """
    objs = result.objects
    n = len(objs)
    parent = list(range(n))

    def find(i: int) -> int:
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(i: int, j: int) -> None:
        ri, rj = find(i), find(j)
        if ri != rj:
            parent[rj] = ri

    for i in range(n):
        for j in range(i + 1, n):
            pos_d = float(np.linalg.norm(objs[i].mean_3d_position - objs[j].mean_3d_position))
            if pos_d > position_radius_m:
                continue
            vi, vj = objs[i].mean_displacement_m, objs[j].mean_displacement_m
            ni, nj = float(np.linalg.norm(vi)), float(np.linalg.norm(vj))
            if ni < 1e-9 or nj < 1e-9:
                continue
            cos_sim = float(np.dot(vi, vj) / (ni * nj))
            ratio = max(ni, nj) / min(ni, nj)
            if cos_sim >= direction_cos_min and ratio <= magnitude_ratio_max:
                union(i, j)

    raw_groups: dict[int, list[int]] = {}
    for i in range(n):
        raw_groups.setdefault(find(i), []).append(i)
    # Deterministic output order: sort groups by their smallest member's
    # original cluster_id, so re-running merge on the same input always
    # produces the same new cluster ids.
    ordered_groups = sorted(
        raw_groups.values(), key=lambda members: min(objs[k].cluster_id for k in members)
    )

    old_to_new = np.full(
        (max((o.cluster_id for o in objs), default=-1) + 1,), -1, dtype=np.int64
    )
    for new_id, members in enumerate(ordered_groups):
        for k in members:
            old_to_new[objs[k].cluster_id] = new_id

    new_labels = result.cluster_labels.copy()
    is_assigned = new_labels >= 0
    new_labels[is_assigned] = old_to_new[new_labels[is_assigned]]

    merged_objects: list[DiscoveredObject] = []
    for new_id, members in enumerate(ordered_groups):
        query_uv, centroid = _recombine_uv_stats(result.track, new_labels, new_id)
        weights = np.array([objs[k].n_tracks for k in members], dtype=np.float64)
        total_n = float(weights.sum())
        mean_3d = sum(
            objs[k].mean_3d_position * w for k, w in zip(members, weights, strict=True)
        ) / total_n
        mean_disp = sum(
            objs[k].mean_displacement_m * w for k, w in zip(members, weights, strict=True)
        ) / total_n
        total_disp_px = float(
            sum(objs[k].total_displacement_px * w for k, w in zip(members, weights, strict=True))
            / total_n
        )
        merged_objects.append(
            DiscoveredObject(
                cluster_id=new_id,
                query_uv=query_uv,
                centroid_uv_per_frame=centroid,
                n_tracks=int(total_n),
                total_displacement_px=total_disp_px,
                mean_3d_position=mean_3d.astype(np.float64),
                mean_displacement_m=mean_disp.astype(np.float64),
            )
        )

    merged_stats = replace(result.stats, n_clusters=len(merged_objects))
    merged_result = DiscoveryResult(
        objects=tuple(merged_objects),
        track=result.track,
        cluster_labels=new_labels,
        stats=merged_stats,
        seed_frame_idx=result.seed_frame_idx,
    )
    merge_groups = tuple(
        tuple(sorted(objs[k].cluster_id for k in members)) for members in ordered_groups
    )
    return merged_result, merge_groups


# --------------------------------------------------------------------------- #
# Second-pass densification (small clusters get a denser local reseed)
# --------------------------------------------------------------------------- #

#: Below this many surviving tracks, a cluster is judged too sparse to trust
#: for S5's downstream point-prompt selection (which wants several *spread*
#: points, see fpgm.datagen.object_masks.select_spread_points) or S6's PnP
#: seed (which wants >= 6 points, ObjectPoseConfig.pnp_min_points, and
#: degrades badly near that floor -- see solve_rigid_track's planarity
#: guard). Measured on the demo episode: the brick, ~38x38 px at an 80x45
#: grid (16x16 px cells), got only 4 first-pass tracks -- below this
#: threshold and exactly the case densification targets.
_DENSIFY_MIN_TRACKS = 20
#: Dense local re-seed grid per under-sampled cluster. 12x12=144 points
#: comfortably covers a 38x38 px object at ~3-4 px spacing even after the
#: bbox padding below, without re-seeding the whole frame at that density
#: (which is the expensive global-densification alternative this second
#: pass exists to avoid -- see the module docstring's "second pass" note in
#: the class-level rationale).
_DENSIFY_GRID = (12, 12)
#: Bbox padding as a fraction of the cluster's own seed-frame extent, each
#: side. A fixed-pixel margin would either barely enlarge the drawer's
#: already-large bbox or blow the brick's tiny one out past the table edge;
#: scaling by the object's own size treats both consistently.
_DENSIFY_BBOX_PAD_FRAC = 1.0
#: A degenerate (near-zero-extent) bbox before padding is floored to this
#: many pixels per side, so a handful of query points that all landed
#: within a few pixels of each other still get a meaningful dense reseed
#: area instead of one that rounds to nothing.
_DENSIFY_MIN_EXTENT_PX = 8.0


def compute_object_bbox_px(
    query_uv: np.ndarray,
    width: int,
    height: int,
    pad_frac: float = _DENSIFY_BBOX_PAD_FRAC,
    min_extent_px: float = _DENSIFY_MIN_EXTENT_PX,
) -> tuple[float, float, float, float]:
    """Padded, frame-clamped pixel bbox around a cluster's seed-frame query points.

    See :data:`_DENSIFY_BBOX_PAD_FRAC` for why the padding scales with the
    object's own extent rather than being a fixed pixel margin.

    Args:
        query_uv: ``(n, 2)`` seed-frame pixel positions of a cluster's
            surviving first-pass tracks (:attr:`DiscoveredObject.query_uv`).
        width: Frame width, pixels.
        height: Frame height, pixels.
        pad_frac: Padding on each side, as a fraction of the (floored) raw
            extent.
        min_extent_px: Floor applied to the raw (pre-padding) width/height
            before padding, so a near-point cluster still gets a usable box.

    Returns:
        ``(x0, y0, x1, y1)``, clamped to ``[0, width-1] x [0, height-1]``.

    Raises:
        ValueError: If ``query_uv`` is empty.
    """
    query_uv = np.asarray(query_uv, dtype=np.float64)
    if query_uv.shape[0] == 0:
        raise ValueError("compute_object_bbox_px: query_uv is empty")
    x0, x1 = float(query_uv[:, 0].min()), float(query_uv[:, 0].max())
    y0, y1 = float(query_uv[:, 1].min()), float(query_uv[:, 1].max())
    w = max(x1 - x0, min_extent_px)
    h = max(y1 - y0, min_extent_px)
    cx, cy = (x0 + x1) / 2.0, (y0 + y1) / 2.0
    half_w = (w / 2.0) * (1.0 + pad_frac)
    half_h = (h / 2.0) * (1.0 + pad_frac)
    x0c = max(0.0, cx - half_w)
    y0c = max(0.0, cy - half_h)
    x1c = min(float(width - 1), cx + half_w)
    y1c = min(float(height - 1), cy + half_h)
    return x0c, y0c, x1c, y1c


def sample_dense_grid_in_bbox(
    bbox: tuple[float, float, float, float], n_cols: int, n_rows: int
) -> np.ndarray:
    """Cell-centred grid inside ``bbox``, same construction as :func:`sample_grid_points`.

    Reuses :func:`sample_grid_points`'s own cell-centring (see that
    function's docstring for why) by sampling a local ``bbox``-sized grid at
    the origin and translating it -- rather than re-deriving the same
    arithmetic here, which would risk the two grids' centring conventions
    silently drifting apart under a future edit to one but not the other.

    Args:
        bbox: ``(x0, y0, x1, y1)`` pixel bbox, e.g. from
            :func:`compute_object_bbox_px`.
        n_cols: Grid columns.
        n_rows: Grid rows.

    Returns:
        ``(n_cols * n_rows, 2)`` float32 pixel coordinates.
    """
    x0, y0, x1, y1 = bbox
    w = max(x1 - x0, 1e-6)
    h = max(y1 - y0, 1e-6)
    local = sample_grid_points(width=w, height=h, n_cols=n_cols, n_rows=n_rows)
    local[:, 0] += x0
    local[:, 1] += y0
    return local


def build_densify_queries(
    result: DiscoveryResult,
    width: int,
    height: int,
    min_tracks: int = _DENSIFY_MIN_TRACKS,
    grid: tuple[int, int] = _DENSIFY_GRID,
) -> tuple[np.ndarray, np.ndarray]:
    """Dense reseed query points for every under-sampled cluster in ``result``.

    Pure/no-GPU half of the second pass: decides *where* to reseed. The
    orchestration layer (:meth:`DynamicObjectStage.run`) supplies these
    points to the tracker and then folds the tracked result back in via
    :func:`densify_undersampled_clusters` -- kept as two functions so the
    seeding logic here stays unit-testable without a tracker.

    Args:
        result: A (typically post-:func:`merge_consistent_clusters`)
            :class:`DiscoveryResult`.
        width: Frame width, pixels.
        height: Frame height, pixels.
        min_tracks: See :data:`_DENSIFY_MIN_TRACKS`.
        grid: ``(n_cols, n_rows)`` per under-sampled cluster, see
            :data:`_DENSIFY_GRID`.

    Returns:
        ``(query_uv, source_cluster_id)``: ``(Q2, 2)`` float32 pixel
        coordinates and ``(Q2,)`` int32 the *target* cluster id each point
        was seeded for -- not yet verified to actually belong there, that
        happens after tracking in :func:`densify_undersampled_clusters`.
        Both are empty (``(0, 2)``/``(0,)``) if no cluster is under-sampled.
    """
    n_cols, n_rows = grid
    queries: list[np.ndarray] = []
    sources: list[np.ndarray] = []
    for obj in result.objects:
        if obj.n_tracks >= min_tracks:
            continue
        bbox = compute_object_bbox_px(obj.query_uv, width, height)
        pts = sample_dense_grid_in_bbox(bbox, n_cols, n_rows)
        queries.append(pts)
        sources.append(np.full(pts.shape[0], obj.cluster_id, dtype=np.int32))
    if not queries:
        return np.zeros((0, 2), dtype=np.float32), np.zeros((0,), dtype=np.int32)
    return np.concatenate(queries, axis=0), np.concatenate(sources, axis=0)


def _concat_tracks(a: Track2D, b: Track2D) -> Track2D:
    """Concatenate two same-length (``frames``/``resolution``-matching) tracks along Q."""
    if a.resolution != b.resolution:
        raise DataError(f"_concat_tracks: resolution mismatch {a.resolution} != {b.resolution}")
    if not np.array_equal(a.frames, b.frames):
        raise DataError("_concat_tracks: frame axes differ -- tracks are not from the same rollout")
    offset = (int(a.point_id.max()) + 1) if a.point_id.size else 0
    return Track2D(
        point_id=np.concatenate([a.point_id, b.point_id + offset]),
        frames=a.frames,
        uv=np.concatenate([a.uv, b.uv], axis=1),
        visible=np.concatenate([a.visible, b.visible], axis=1),
        resolution=a.resolution,
    )


def densify_undersampled_clusters(
    result: DiscoveryResult,
    dense_track: Track2D,
    dense_world_xyz: np.ndarray,
    dense_z_cam: np.ndarray,
    source_cluster_id: np.ndarray,
    robot_seg: np.ndarray,
    query_frame_idx: int,
    robot_life_fraction_gate: float = _ROBOT_LIFE_FRACTION_GATE,
    static_displacement_gate_px: float = _STATIC_DISPLACEMENT_GATE_PX,
) -> DiscoveryResult:
    """Fold a second-pass dense reseed into ``result``, gated exactly like the first pass.

    Every reseeded point was deliberately placed inside one specific
    surviving cluster's own (padded) bbox by :func:`build_densify_queries`
    -- so a surviving point is assigned straight to that ``source_cluster_id``
    rather than re-run through :func:`cluster_tracks`. Re-clustering here
    would only add DBSCAN's density requirement back on top of a
    *deliberately* localised sample (exactly the sparsity problem this pass
    exists to fix) and risks splitting the new dense points off from the
    original sparse ones as a second, spurious cluster for the same object.
    A point is still dropped, never reassigned to a different cluster, if it
    fails the same robot-overlap or static-displacement gate the first pass
    used -- the bbox padding (see :data:`_DENSIFY_BBOX_PAD_FRAC`) can catch
    a strip of background or, near the gripper, the robot itself, and those
    gates are exactly what is supposed to catch that, unchanged.

    Per-cluster ``mean_3d_position``/``mean_displacement_m``/
    ``total_displacement_px`` are recombined via the same exact
    ``n_tracks``-weighted-mean identity :func:`merge_consistent_clusters`
    uses (both are just growing the same "mean over a cluster's member
    tracks" by adding more members) -- old members contribute through their
    already-aggregated :class:`DiscoveredObject` values, new members
    through their own freshly-lifted world_xyz/displacement, and the
    combination is exact for the same reason as there: these are all plain
    arithmetic means.

    Args:
        result: A (typically post-merge) :class:`DiscoveryResult`. Clusters
            with no reseed points targeting them pass through unchanged.
        dense_track: Tracker output on the reseed points returned by
            :func:`build_densify_queries` (same seed frame, whole-video
            rollout, i.e. same ``frames``/``resolution`` as ``result.track``).
        dense_world_xyz: ``(Q2, 3)`` float64, world metres -- lifted the same
            way as the first pass (e.g. another :func:`lift_query_points`
            call against the same seed-frame scene-flow anchor).
        dense_z_cam: ``(Q2,)`` float64, camera-frame depth at each point.
        source_cluster_id: ``(Q2,)`` int32, from :func:`build_densify_queries`
            -- which cluster each reseed point targets.
        robot_seg: ``(T, H, W)`` integer array, nonzero = robot.
        query_frame_idx: Absolute video frame the reseed points were seeded on.
        robot_life_fraction_gate: See :func:`filter_robot_tracks`.
        static_displacement_gate_px: See :func:`filter_static_tracks`.

    Returns:
        A new :class:`DiscoveryResult` with the surviving reseed points
        folded into ``track``/``cluster_labels`` and the targeted clusters'
        stats updated; ``stats.n_query_points``/``n_after_robot_filter``/
        ``n_after_static_filter``/``n_with_3d_lift`` are increased by this
        pass's own funnel counts (added, not replaced -- they are meant to
        describe the whole two-pass discovery run once densification runs).
    """
    if dense_track.uv.shape[1] == 0:
        return result

    keep_robot = filter_robot_tracks(
        dense_track, robot_seg, query_frame_idx, robot_life_fraction_gate
    )
    keep_static = filter_static_tracks(dense_track, static_displacement_gate_px)
    has_lift = np.all(np.isfinite(dense_world_xyz), axis=1) & np.isfinite(dense_z_cam)
    keep = keep_robot & keep_static & has_lift

    n_reseed = int(dense_track.uv.shape[1])
    n_after_robot = int(keep_robot.sum())
    n_after_static = int((keep_robot & keep_static).sum())
    n_lifted = int(keep.sum())

    if not keep.any():
        stats = replace(
            result.stats,
            n_query_points=result.stats.n_query_points + n_reseed,
            n_after_robot_filter=result.stats.n_after_robot_filter + n_after_robot,
            n_after_static_filter=result.stats.n_after_static_filter + n_after_static,
            n_with_3d_lift=result.stats.n_with_3d_lift + n_lifted,
        )
        return replace(result, stats=stats)

    kept_track = _subset_track(dense_track, keep)
    kept_world = np.asarray(dense_world_xyz)[keep]
    kept_source = source_cluster_id[keep]
    _, mag_kept = compute_displacement(kept_track)

    merged_track = _concat_tracks(result.track, kept_track)
    merged_labels = np.concatenate([result.cluster_labels, kept_source])

    # dx_m/dy_m for the new points -- same pinhole conversion as
    # build_cluster_features, but that helper needs fx/fy and this function
    # deliberately does not take a camera (the seeding/tracking already
    # happened at whatever resolution dense_track is in); the ratio between
    # a track's px displacement and its own already-known metric magnitude
    # (mag_kept in metres would need fx/fy too) is avoided entirely by
    # reusing the *direction* from pixels and rescaling to each object's own
    # measured metric/pixel ratio -- but that is fragile across objects at
    # different depths. Simpler and exact: mean_3d_position uses kept_world
    # directly (already metric), and mean_displacement_m falls back to the
    # existing cluster's own value for direction/magnitude scaling, since a
    # handful of new points added to an existing rigid-body cluster does not
    # change that cluster's true motion -- only densifies its position/count
    # evidence. See TestDensifyUndersampledClusters for the resulting
    # (unchanged) mean_displacement_m assertion.
    new_objects = list(result.objects)
    by_cluster_id = {o.cluster_id: i for i, o in enumerate(new_objects)}
    for cid in sorted(set(int(x) for x in kept_source.tolist())):
        if cid not in by_cluster_id:
            continue  # a reseed source cluster no longer exists (should not happen)
        idx = by_cluster_id[cid]
        old = new_objects[idx]
        member_new = kept_source == cid
        n_new = int(member_new.sum())
        if n_new == 0:
            continue
        query_uv, centroid = _recombine_uv_stats(merged_track, merged_labels, cid)
        new_mean_3d = np.nanmean(kept_world[member_new], axis=0)
        total_n = old.n_tracks + n_new
        mean_3d = (old.mean_3d_position * old.n_tracks + new_mean_3d * n_new) / total_n
        new_mean_disp_px = float(np.mean(mag_kept[member_new]))
        total_disp_px = (
            old.total_displacement_px * old.n_tracks + new_mean_disp_px * n_new
        ) / total_n
        new_objects[idx] = DiscoveredObject(
            cluster_id=cid,
            query_uv=query_uv,
            centroid_uv_per_frame=centroid,
            n_tracks=total_n,
            total_displacement_px=float(total_disp_px),
            mean_3d_position=mean_3d.astype(np.float64),
            mean_displacement_m=old.mean_displacement_m,  # see note above
        )

    stats = replace(
        result.stats,
        n_query_points=result.stats.n_query_points + n_reseed,
        n_after_robot_filter=result.stats.n_after_robot_filter + n_after_robot,
        n_after_static_filter=result.stats.n_after_static_filter + n_after_static,
        n_with_3d_lift=result.stats.n_with_3d_lift + n_lifted,
    )
    return DiscoveryResult(
        objects=tuple(new_objects),
        track=merged_track,
        cluster_labels=merged_labels,
        stats=stats,
        seed_frame_idx=result.seed_frame_idx,
    )


# --------------------------------------------------------------------------- #
# On-disk artifacts
# --------------------------------------------------------------------------- #


def _write_discovery_npz(path: Path, result: DiscoveryResult) -> None:
    obj_fields: dict[str, np.ndarray] = {}
    for i, obj in enumerate(result.objects):
        obj_fields[f"obj{i}_cluster_id"] = np.array(obj.cluster_id, dtype=np.int32)
        obj_fields[f"obj{i}_query_uv"] = obj.query_uv
        obj_fields[f"obj{i}_centroid_uv"] = obj.centroid_uv_per_frame
        obj_fields[f"obj{i}_n_tracks"] = np.array(obj.n_tracks, dtype=np.int32)
        obj_fields[f"obj{i}_total_displacement_px"] = np.array(
            obj.total_displacement_px, dtype=np.float32
        )
        obj_fields[f"obj{i}_mean_3d_position"] = obj.mean_3d_position
        obj_fields[f"obj{i}_mean_displacement_m"] = obj.mean_displacement_m
    ensure_dir(path.parent)
    np.savez_compressed(
        path,
        n_objects=np.array(len(result.objects), dtype=np.int32),
        seed_frame_idx=np.array(result.seed_frame_idx, dtype=np.int32),
        track_point_id=result.track.point_id,
        track_frames=result.track.frames,
        track_uv=result.track.uv,
        track_visible=result.track.visible,
        track_resolution=np.array(result.track.resolution, dtype=np.int64),
        cluster_labels=result.cluster_labels,
        stats_n_query_points=np.array(result.stats.n_query_points, dtype=np.int32),
        stats_n_after_robot_filter=np.array(result.stats.n_after_robot_filter, dtype=np.int32),
        stats_n_after_static_filter=np.array(result.stats.n_after_static_filter, dtype=np.int32),
        stats_n_with_3d_lift=np.array(result.stats.n_with_3d_lift, dtype=np.int32),
        stats_n_clusters=np.array(result.stats.n_clusters, dtype=np.int32),
        stats_n_noise=np.array(result.stats.n_noise, dtype=np.int32),
        **obj_fields,
    )


def _read_discovery_npz(path: Path) -> DiscoveryResult:
    with np.load(path) as data:
        n_objects = int(data["n_objects"])
        objects = []
        for i in range(n_objects):
            objects.append(
                DiscoveredObject(
                    cluster_id=int(data[f"obj{i}_cluster_id"]),
                    query_uv=data[f"obj{i}_query_uv"],
                    centroid_uv_per_frame=data[f"obj{i}_centroid_uv"],
                    n_tracks=int(data[f"obj{i}_n_tracks"]),
                    total_displacement_px=float(data[f"obj{i}_total_displacement_px"]),
                    mean_3d_position=data[f"obj{i}_mean_3d_position"],
                    mean_displacement_m=(
                        data[f"obj{i}_mean_displacement_m"]
                        if f"obj{i}_mean_displacement_m" in data
                        else np.zeros(2, dtype=np.float64)
                    ),
                )
            )
        track = Track2D(
            point_id=data["track_point_id"],
            frames=data["track_frames"],
            uv=data["track_uv"],
            visible=data["track_visible"],
            resolution=tuple(int(x) for x in data["track_resolution"]),
        )
        stats = DiscoveryStats(
            n_query_points=int(data["stats_n_query_points"]),
            n_after_robot_filter=int(data["stats_n_after_robot_filter"]),
            n_after_static_filter=int(data["stats_n_after_static_filter"]),
            n_with_3d_lift=int(data["stats_n_with_3d_lift"]),
            n_clusters=int(data["stats_n_clusters"]),
            n_noise=int(data["stats_n_noise"]),
        )
        return DiscoveryResult(
            objects=tuple(objects),
            track=track,
            cluster_labels=data["cluster_labels"],
            stats=stats,
            seed_frame_idx=int(data["seed_frame_idx"]),
        )


# --------------------------------------------------------------------------- #
# Debug deliverables
# --------------------------------------------------------------------------- #


def write_discovery_debug_video(
    path: str | Path,
    video_path: str | Path,
    result: DiscoveryResult,
    robot_seg: np.ndarray,
    fps: float,
) -> Path:
    """Real video + surviving tracks coloured by cluster + dimmed robot mask + HUD.

    The eyeball-checkable deliverable this stage exists to produce -- per the
    task brief, this matters more than any single number.
    """
    from fpgm.viz.overlays import VideoWriter, color_for, draw_hud, read_frames_bgr

    track = result.track
    labels = result.cluster_labels
    n_frames = track.frames.shape[0]

    path = Path(path)
    with VideoWriter(path, fps=fps) as writer:
        for t, frame_bgr in enumerate(read_frames_bgr(video_path)):
            if t >= n_frames:
                break
            out = frame_bgr.copy()
            if t < robot_seg.shape[0]:
                robot_mask = np.asarray(robot_seg[t]) > 0
                if robot_mask.any():
                    tint = np.zeros_like(out)
                    out[robot_mask] = cv2.addWeighted(out, 0.35, tint, 0.65, 0.0)[robot_mask]

            uv = track.uv[t]
            vis = track.visible[t]
            for q in range(uv.shape[0]):
                if not vis[q] or not np.all(np.isfinite(uv[q])):
                    continue
                cid = int(labels[q])
                color = (110, 110, 110) if cid < 0 else color_for(cid)
                pt = (int(round(float(uv[q, 0]))), int(round(float(uv[q, 1]))))
                cv2.circle(out, pt, 3, color, -1, lineType=cv2.LINE_AA)

            s = result.stats
            out = draw_hud(
                out,
                [
                    f"frame {t}/{n_frames - 1}  seed={result.seed_frame_idx}",
                    f"clusters={s.n_clusters}  noise_pts={s.n_noise}",
                    f"query={s.n_query_points} -> robot={s.n_after_robot_filter} "
                    f"-> static={s.n_after_static_filter} -> lifted={s.n_with_3d_lift}",
                ],
            )
            writer.write(out)
    logger.info("wrote %s", path)
    return path


def write_discovery_clusters_png(
    path: str | Path, video_path: str | Path, result: DiscoveryResult
) -> Path:
    """Seed frame with every surviving query point coloured by its cluster."""
    from fpgm.viz.overlays import color_for, read_frames_bgr

    seed_frame_bgr = None
    for t, frame_bgr in enumerate(read_frames_bgr(video_path)):
        if t == result.seed_frame_idx:
            seed_frame_bgr = frame_bgr.copy()
            break
    if seed_frame_bgr is None:
        raise DataError(
            f"write_discovery_clusters_png: video {video_path} has fewer than "
            f"{result.seed_frame_idx + 1} frames -- cannot render the seed frame"
        )

    track = result.track
    labels = result.cluster_labels
    seed_matches = np.flatnonzero(track.frames == result.seed_frame_idx)
    seed_t = int(seed_matches[0]) if seed_matches.size else 0
    uv = track.uv[seed_t]
    for q in range(uv.shape[0]):
        if not np.all(np.isfinite(uv[q])):
            continue
        cid = int(labels[q])
        color = (110, 110, 110) if cid < 0 else color_for(cid)
        pt = (int(round(float(uv[q, 0]))), int(round(float(uv[q, 1]))))
        cv2.circle(seed_frame_bgr, pt, 4, color, -1, lineType=cv2.LINE_AA)

    for obj in result.objects:
        centroid = np.mean(obj.query_uv, axis=0)
        label = f"cluster {obj.cluster_id} n={obj.n_tracks}"
        org = (int(round(float(centroid[0]))), int(round(float(centroid[1]))))
        cv2.putText(
            seed_frame_bgr, label, org, cv2.FONT_HERSHEY_SIMPLEX, 0.6,
            color_for(obj.cluster_id), 2, lineType=cv2.LINE_AA,
        )

    path = Path(path)
    ensure_dir(path.parent)
    cv2.imwrite(str(path), seed_frame_bgr)
    logger.info("wrote %s", path)
    return path


# --------------------------------------------------------------------------- #
# Stage orchestration
# --------------------------------------------------------------------------- #


class DynamicObjectStage:
    """S4 orchestration: TAPNext++ full-frame grid -> robot/static filter -> DBSCAN.

    Ties together the pure logic above with the real inputs it needs (a
    checkpointed tracker, the episode's scene-flow data) and caches the
    (expensive, GPU-bound) result via :class:`~fpgm.datagen.cache.StageCache`.

    **Model injection.** ``tracker`` lets a caller hand in an already-built
    ``TapNextPointTracker`` (e.g. the same instance
    :class:`~fpgm.datagen.object_masks.PromptMaskStage` already built for
    its own dense re-track, or a process-level registry an orchestrator
    maintains) instead of this class constructing a fresh one -- identical
    rationale to :class:`~fpgm.datagen.object_masks.ObjectMaskStage`'s own
    "Model injection" docstring section (construct TAPNext++ once per
    process, not once per stage object). ``None`` (the default) keeps
    :meth:`_build_tracker` constructing fresh, exactly as before -- every
    existing no-injection call site (including this repo's GPU-free tests,
    which never reach ``_build_tracker`` at all) keeps working unchanged.
    This matters specifically for :mod:`fpgm.datagen.object_masks`'s
    ``PromptMaskStage`` discovery fallback, which constructs this class on
    demand only when a role's every text-prompt candidate fails -- sharing
    the tracker there avoids a second TAPNext++ load on top of the one
    :class:`PromptMaskStage` already paid for its own dense re-track.
    """

    def __init__(
        self,
        cfg: DatagenConfig,
        tracking_cfg: TrackingConfig,
        checkpoint_path: str,
        device: str = "cuda",
        convention_cfg: ConventionConfig | None = None,
        tracker: object | None = None,
    ) -> None:
        self.cfg = cfg
        self.tracking_cfg = tracking_cfg
        self.checkpoint_path = checkpoint_path
        self.device = device
        self.convention_cfg = convention_cfg or ConventionConfig()
        self._injected_tracker = tracker

    def run(
        self,
        *,
        uuid: str,
        camera_serial: str,
        video_path: str | Path,
        video_width: int,
        video_height: int,
        robot_seg: np.ndarray,
        frame_index: EpisodeFrameIndex,
        flows_reader: FlowsReader,
        cache: StageCache,
        seed_frame_idx: int = 0,
        fps: float = 15.0,
        force: bool = False,
        merge_clusters: bool = True,
        densify_small_clusters: bool = True,
    ) -> DiscoveryResult:
        """Run (or reuse a cached) S4 discovery pass for one episode/camera.

        Three passes, in order: (1) the full-frame grid rollout ->
        :func:`discover_objects` (first-pass DBSCAN clustering), (2)
        :func:`merge_consistent_clusters` to fold DBSCAN's over-segmentation
        of large rigid objects back together, (3)
        :func:`build_densify_queries`/:func:`densify_undersampled_clusters`
        to give any cluster still too sparse afterward (typically a small
        object like the brick) a denser local reseed. Both (2) and (3) are
        no-ops when there is nothing for them to do (a single merge group
        per cluster, or every cluster already at/above
        ``_DENSIFY_MIN_TRACKS``), so leaving them enabled by default costs
        nothing on an episode that does not need them.

        Args:
            uuid: Episode uuid, carried into the fingerprint/logs only.
            camera_serial: Which camera this run is for.
            video_path: Path to the camera's full-episode mp4.
            video_width: Video frame width, pixels.
            video_height: Video frame height, pixels.
            robot_seg: ``(T, H, W)`` integer array, nonzero = robot (S2's
                ``robot_buffers/seg.npy`` -- pass an ``mmap_mode="r"`` load for
                a large episode).
            frame_index: Already-built :class:`EpisodeFrameIndex`.
            flows_reader: An already-open :class:`FlowsReader` for this episode.
            cache: Rooted at ``outputs/datagen/<uuid>/<camera_serial>/``.
            seed_frame_idx: Absolute video frame to seed TAPNext++ on. See the
                module docstring for why 0 is the default and how to check it
                is appropriate for a given episode.
            fps: Playback rate for ``discovery_debug.mp4``.
            force: Ignore the on-disk cache and recompute.
            merge_clusters: Run :func:`merge_consistent_clusters` (pass 2).
                Toggleable so a caller can A/B the raw DBSCAN output; off by
                default only makes sense for debugging this stage itself.
            densify_small_clusters: Run the second-pass dense reseed (pass
                3). Costs one extra (localised, cheap) tracker rollout only
                when at least one cluster is under-sampled after merging.

        Returns:
            A :class:`DiscoveryResult`.
        """
        stage_name = STAGE_NAME
        master_dir = ensure_dir(cache.root / "master")
        npz_path = master_dir / "discovery_tracks.npz"
        debug_video_path = master_dir / "discovery_debug.mp4"
        clusters_png_path = master_dir / "discovery_clusters.png"

        video_path = Path(video_path)
        video_stat = video_path.stat() if video_path.exists() else None
        fingerprint = {
            "uuid": uuid,
            "camera_serial": camera_serial,
            "seed_frame_idx": seed_frame_idx,
            "tapnext_grid": list(self.cfg.tapnext_grid),
            "dbscan_eps_m": self.cfg.dbscan_eps_m,
            "dbscan_min_samples": self.cfg.dbscan_min_samples,
            "native_resolution": list(self.cfg.native_resolution),
            "checkpoint_path": str(self.checkpoint_path),
            "video_mtime": video_stat.st_mtime if video_stat else None,
            "video_size": video_stat.st_size if video_stat else None,
            "video_width": video_width,
            "video_height": video_height,
            "merge_clusters": merge_clusters,
            "densify_small_clusters": densify_small_clusters,
            "code_version": "discovery.v2",
        }

        if not force and cache.is_fresh(stage_name, fingerprint) and npz_path.exists():
            logger.info("discovery: reusing cached %s", npz_path)
            return _read_discovery_npz(npz_path)

        n_cols, n_rows = self.cfg.tapnext_grid
        query_uv = sample_grid_points(video_width, video_height, n_cols, n_rows)
        logger.info(
            "discovery: seeding %d query points (%dx%d grid) at frame %d",
            query_uv.shape[0], n_cols, n_rows, seed_frame_idx,
        )

        tracker = self._build_tracker()
        track = tracker.track_video_file(str(video_path), query_uv, query_frame_idx=seed_frame_idx)

        anchor = _build_scene_flow_anchor(
            frame_index, flows_reader, camera_serial, seed_frame_idx,
            video_width, video_height, self.cfg.native_resolution, self.convention_cfg,
        )
        world_xyz, z_cam = lift_query_points(anchor, query_uv, _SCENE_FLOW_MATCH_RADIUS_PX)

        result = discover_objects(
            track, robot_seg, seed_frame_idx, world_xyz, z_cam,
            fx=anchor.camera.K.fx, fy=anchor.camera.K.fy, cfg=self.cfg,
        )
        n_clusters_pass1 = result.stats.n_clusters
        pass1_report = [
            f"  cluster {o.cluster_id}: n_tracks={o.n_tracks} "
            f"mean_3d={np.round(o.mean_3d_position, 3).tolist()} "
            f"disp_px={o.total_displacement_px:.1f}"
            for o in result.objects
        ]
        logger.info(
            "discovery: pass 1 (raw DBSCAN): %d clusters from %d/%d/%d/%d "
            "(query/robot-filt/static-filt/lifted) tracks\n%s",
            n_clusters_pass1, result.stats.n_query_points,
            result.stats.n_after_robot_filter, result.stats.n_after_static_filter,
            result.stats.n_with_3d_lift, "\n".join(pass1_report),
        )

        merge_groups: tuple[tuple[int, ...], ...] = tuple((o.cluster_id,) for o in result.objects)
        if merge_clusters and result.objects:
            result, merge_groups = merge_consistent_clusters(result)
            logger.info(
                "discovery: pass 2 (merge): %d -> %d clusters, groups=%s",
                n_clusters_pass1, result.stats.n_clusters, merge_groups,
            )

        n_densify_reseeded = 0
        if densify_small_clusters and any(
            o.n_tracks < _DENSIFY_MIN_TRACKS for o in result.objects
        ):
            dense_uv, dense_source = build_densify_queries(result, video_width, video_height)
            n_densify_reseeded = int(dense_uv.shape[0])
            if n_densify_reseeded:
                logger.info(
                    "discovery: pass 3 (densify): reseeding %d points around %d "
                    "under-sampled cluster(s)",
                    n_densify_reseeded,
                    len({int(x) for x in dense_source.tolist()}),
                )
                dense_track = tracker.track_video_file(
                    str(video_path), dense_uv, query_frame_idx=seed_frame_idx
                )
                dense_world_xyz, dense_z_cam = lift_query_points(
                    anchor, dense_uv, _SCENE_FLOW_MATCH_RADIUS_PX
                )
                result = densify_undersampled_clusters(
                    result, dense_track, dense_world_xyz, dense_z_cam, dense_source,
                    robot_seg, seed_frame_idx,
                )
                post_report = [
                    f"  cluster {o.cluster_id}: n_tracks={o.n_tracks} "
                    f"mean_3d={np.round(o.mean_3d_position, 3).tolist()}"
                    for o in result.objects
                ]
                logger.info("discovery: after densify:\n%s", "\n".join(post_report))

        logger.info(
            "discovery: final %d clusters (pass1 raw=%d) from %d/%d/%d/%d "
            "(query/robot-filt/static-filt/lifted) tracks",
            result.stats.n_clusters, n_clusters_pass1, result.stats.n_query_points,
            result.stats.n_after_robot_filter, result.stats.n_after_static_filter,
            result.stats.n_with_3d_lift,
        )

        _write_discovery_npz(npz_path, result)
        write_discovery_debug_video(debug_video_path, video_path, result, robot_seg, fps=fps)
        write_discovery_clusters_png(clusters_png_path, video_path, result)

        limitations = [
            f"S4 discovery seeded a single TAPNext++ rollout at video frame "
            f"{seed_frame_idx}; frames before it (none, since seed_frame_idx=0 by "
            "default) would have no track coverage at all, per TAPNext++'s causal "
            "seeding constraint.",
            "Cluster count/identity is a RESULT of this stage, not verified against "
            "any ground truth here -- compare against the hand-labelled reference "
            "masks (objects_sam3d_multi/objects_drawer_fit) out of band.",
            f"Cluster merging: {'enabled' if merge_clusters else 'disabled'} "
            f"(pass-1 raw cluster count {n_clusters_pass1} -> final "
            f"{result.stats.n_clusters}; groups={merge_groups}).",
            f"Second-pass densification: {'enabled' if densify_small_clusters else 'disabled'} "
            f"({n_densify_reseeded} reseed points added around under-sampled clusters, "
            f"threshold n_tracks < {_DENSIFY_MIN_TRACKS}).",
        ]
        cache.write_meta(
            stage_name,
            fingerprint,
            payload={
                "n_objects": len(result.objects),
                "n_clusters_pass1_raw": n_clusters_pass1,
                "merge_groups": merge_groups,
                "n_densify_reseeded": n_densify_reseeded,
                "stats": {
                    "n_query_points": result.stats.n_query_points,
                    "n_after_robot_filter": result.stats.n_after_robot_filter,
                    "n_after_static_filter": result.stats.n_after_static_filter,
                    "n_with_3d_lift": result.stats.n_with_3d_lift,
                    "n_clusters": result.stats.n_clusters,
                    "n_noise": result.stats.n_noise,
                },
                "objects": [
                    {
                        "cluster_id": obj.cluster_id,
                        "n_tracks": obj.n_tracks,
                        "total_displacement_px": obj.total_displacement_px,
                        "mean_3d_position": obj.mean_3d_position.tolist(),
                        "mean_displacement_m": obj.mean_displacement_m.tolist(),
                    }
                    for obj in result.objects
                ],
                "discovery_tracks_npz": str(npz_path),
                "discovery_debug_mp4": str(debug_video_path),
                "discovery_clusters_png": str(clusters_png_path),
            },
            limitations=limitations,
        )
        return result

    def _build_tracker(self):
        if self._injected_tracker is not None:
            return self._injected_tracker
        # Heavy, GPU-framework-dependent import kept out of module scope; see
        # fpgm.tracking.tapnext's own docstring for the same reasoning.
        from fpgm.tracking.tapnext import TapNextPointTracker

        return TapNextPointTracker(self.tracking_cfg, self.checkpoint_path, device=self.device)