"""Batched GPU furthest-point sampling for the stage-2 pipeline. Designed for the v2 (FPS) stage-2 path where the dataloader returns the full hull-filtered point cloud padded to a fixed cap, and we subsample to the training budget (e.g. 8k+8k) *after* collate — i.e. once per batch on the GPU inside the training/inference loop. CPU dataloader workers cannot use CUDA, so this is the canonical place to run FPS. Two entry points: * :func:`fps_batched_indices` — single-provenance FPS. Returns ``(B, k)`` long indices into ``(B, N, 3)`` xyz. Honors a per-point ``valid_mask`` (padding slots are never selected) and an optional per-sample RNG seed. * :func:`fps_subsample_stage2_batch` — convenience wrapper that operates on a stage-2 batch dict produced by ``stage2_collate_fn`` with the FPS variant. Runs FPS once per provenance (COLMAP, depth) and gathers all the per-point fields (xyz, semantics, RGB, …) to the resampled order. Returns a NEW batch dict with the same keys but with the per-point tensors shrunk from ``N`` to ``n_col + n_dep`` and ``scene_xyz`` re-normalized into the post-subsample bbox (matching the random-sampling path). The kernel is a vectorized version of the single-loop torch FPS in ``viz_stage2_density._fps_indices``: one ``(B, N)`` distance buffer, ``k`` iterations, each step gathers the just-picked point and updates the running minimum-distance buffer in place. No host syncs inside the loop. """ from __future__ import annotations from typing import Optional import torch # Sentinel for "this slot is ineligible / already picked." Using -inf so the # argmax never selects it, while leaving real (non-negative) squared distances # untouched by the min(). def _neg_inf(dtype: torch.dtype) -> float: return torch.finfo(dtype).min @torch.no_grad() def fps_batched_indices( xyz: torch.Tensor, valid_mask: torch.Tensor, k: int, seed_per_sample: Optional[torch.Tensor] = None, max_exact_iters: Optional[int] = None, ) -> torch.Tensor: """Batched furthest-point sampling on GPU (or CPU; CUDA strongly preferred). Args: xyz: ``(B, N, 3)`` float tensor of point coordinates. valid_mask: ``(B, N)`` bool tensor. ``False`` slots are padding / wrong-provenance / otherwise ineligible. They are never selected. k: number of points to pick per batch element. seed_per_sample: optional ``(B,)`` long tensor. Each entry is taken modulo the per-sample valid count to pick the seed within the eligible set. ``None`` → seed at the first eligible index (deterministic). max_exact_iters: optional cap for exact FPS iterations. If set and smaller than ``k``, exact FPS is used for the first anchors and the remainder is filled with deterministic random eligible points. This keeps broad spatial anchors while avoiding thousands of sequential distance-buffer updates in the stage-2 training hot path. Returns: ``(B, k)`` long tensor of indices into the ``N`` axis. For batch elements with fewer than ``k`` eligible points, the deficit slots are padded by repeating earlier picks (FPS naturally cycles when the eligible set is exhausted and argmax falls back to index 0; the downstream encoder doesn't care about duplicates). """ if xyz.dim() != 3 or xyz.shape[-1] != 3: raise ValueError(f"xyz must be (B, N, 3), got {tuple(xyz.shape)}") if valid_mask.shape != xyz.shape[:2]: raise ValueError( f"valid_mask shape {tuple(valid_mask.shape)} != xyz[:2] " f"{tuple(xyz.shape[:2])}" ) B, N, _ = xyz.shape device = xyz.device if k <= 0: return torch.empty((B, 0), device=device, dtype=torch.long) k_exact = k if max_exact_iters is not None: k_exact = max(0, min(k, int(max_exact_iters))) if k_exact == 0: return _deterministic_random_fill_indices( valid_mask, k=k, seed_per_sample=seed_per_sample, exclude_idx=None, ) NEG = _neg_inf(xyz.dtype) dist2 = torch.where( valid_mask, torch.full_like(xyz[..., 0], float("inf")), torch.full_like(xyz[..., 0], NEG), ) # --- Seed selection ----------------------------------------------------- if seed_per_sample is None: # First eligible index (argmax over bool returns index of first True). seed_idx = valid_mask.to(torch.int8).argmax(dim=1) else: if seed_per_sample.shape != (B,): raise ValueError( f"seed_per_sample shape {tuple(seed_per_sample.shape)} != ({B},)" ) counts = valid_mask.sum(dim=1).clamp(min=1) offset = (seed_per_sample.to(device=device, dtype=torch.long).abs() % counts) # Find the (offset+1)-th True in valid_mask along dim=1. cum = valid_mask.to(torch.long).cumsum(dim=1) target = (offset + 1).unsqueeze(1) seed_idx = (cum >= target).to(torch.int8).argmax(dim=1) # Batches with zero eligible points end up with seed_idx=0; that's # fine — they'll just emit duplicate index 0s and the caller's # valid_mask handling deals with it. arange_b = torch.arange(B, device=device) out = torch.empty((B, k_exact), device=device, dtype=torch.long) out[:, 0] = seed_idx # Mark seed as picked so it can't be re-selected. dist2[arange_b, seed_idx] = NEG last_idx = seed_idx for i in range(1, k_exact): last_pts = xyz[arange_b, last_idx] # (B, 3) diff = xyz - last_pts.unsqueeze(1) # (B, N, 3) new_d2 = (diff * diff).sum(dim=-1) # (B, N) # Only update eligible (still-in-the-game) slots. Picked/ineligible # slots sit at NEG and must stay there. eligible = dist2 > NEG dist2 = torch.where(eligible, torch.minimum(dist2, new_d2), dist2) next_idx = dist2.argmax(dim=1) out[:, i] = next_idx dist2[arange_b, next_idx] = NEG last_idx = next_idx if k_exact < k: fill = _deterministic_random_fill_indices( valid_mask, k=k - k_exact, seed_per_sample=seed_per_sample, exclude_idx=out, ) out = torch.cat([out, fill], dim=1) return out @torch.no_grad() def _deterministic_random_fill_indices( valid_mask: torch.Tensor, k: int, seed_per_sample: Optional[torch.Tensor], exclude_idx: Optional[torch.Tensor], ) -> torch.Tensor: """Pick deterministic pseudo-random valid indices, excluding anchors. This is intentionally one vectorized top-k instead of a loop. It is used only for the fast hybrid FPS mode, where the expensive exact FPS anchors have already provided spatial coverage. """ B, N = valid_mask.shape device = valid_mask.device if k <= 0: return torch.empty((B, 0), device=device, dtype=torch.long) remaining = valid_mask.clone() if exclude_idx is not None and exclude_idx.numel() > 0: remaining.scatter_(1, exclude_idx, False) n_top = min(k, N) idx_f = torch.arange(N, device=device, dtype=torch.float32).unsqueeze(0) batch_f = torch.arange(B, device=device, dtype=torch.float32).unsqueeze(1) if seed_per_sample is None: seed_f = torch.zeros((B, 1), device=device, dtype=torch.float32) else: seed_i = seed_per_sample.to(device=device, dtype=torch.long).remainder(1_000_003) seed_f = seed_i.to(dtype=torch.float32).view(B, 1) # A small deterministic hash in float space. Quality only needs to be good # enough to decorrelate the non-anchor fill; exact FPS has already handled # the coverage-critical prefix. scores = torch.sin((idx_f + 1.0) * 12.9898 + seed_f * 78.233 + batch_f * 37.719) scores = scores * 43758.5453 scores = scores - torch.floor(scores) scores = torch.where(remaining, scores, torch.full_like(scores, -1.0)) fill = scores.topk(n_top, dim=1).indices fill_ok = remaining.gather(1, fill) first_valid = valid_mask.to(torch.int8).argmax(dim=1) if exclude_idx is not None and exclude_idx.numel() > 0: fallback_src = exclude_idx else: fallback_src = first_valid.unsqueeze(1) fallback = fallback_src[:, torch.arange(n_top, device=device) % fallback_src.shape[1]] fill = torch.where(fill_ok, fill, fallback) if n_top < k: pad = fill[:, torch.arange(k - n_top, device=device) % n_top] fill = torch.cat([fill, pad], dim=1) return fill @torch.no_grad() def fps_per_provenance_indices( xyz: torch.Tensor, type_ids: torch.Tensor, valid_mask: torch.Tensor, n_col: int, n_dep: int, seed_per_sample: Optional[torch.Tensor] = None, max_exact_iters_per_provenance: Optional[int] = None, ) -> torch.Tensor: """Per-provenance FPS: COLMAP (type_id==0) first, then depth (type_id==1). Returns ``(B, n_col + n_dep)`` indices into ``N``. The two halves are concatenated so the encoder sees the same provenance ordering as the random-sampling path. """ col_mask = valid_mask & (type_ids == 0) dep_mask = valid_mask & (type_ids == 1) # Use independent seeds so the two provenances don't pick correlated # starting points across the batch. seed_col = seed_per_sample seed_dep = (None if seed_per_sample is None else seed_per_sample ^ 0x9E3779B97F4A7C15) idx_col = fps_batched_indices( xyz, col_mask, n_col, seed_col, max_exact_iters=max_exact_iters_per_provenance, ) idx_dep = fps_batched_indices( xyz, dep_mask, n_dep, seed_dep, max_exact_iters=max_exact_iters_per_provenance, ) return torch.cat([idx_col, idx_dep], dim=1) # --------------------------------------------------------------------------- # Stage-2 batch convenience: gather + re-normalize # --------------------------------------------------------------------------- # Per-point tensors in the stage-2 batch. Matches STAGE2_POINT_KEYS in # data.stage2_build, with the "raw_" prefix dropped by stage2_row_to_sample # and remapped to the "scene_*" namespace. _STAGE2_SCENE_PER_POINT_KEYS = ( "scene_xyz", "scene_type_ids", "scene_gestalt_ids", "scene_gestalt_id2", "scene_gestalt_w1", "scene_ade_ids", "scene_geom_conf", "scene_sem_conf", "scene_rgb", ) def _robust_norm_params_batched( xyz: torch.Tensor, valid_mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Per-batch (median centre, P95 scale) computed only over valid points. Mirrors :func:`data.stage2_dataset._robust_norm_params` for a batched tensor. Implements median + P95 via :func:`torch.quantile` on padded rows; invalid slots are set to NaN so quantile ignores them. """ B = xyz.shape[0] device = xyz.device # Set invalid slots to NaN so torch.quantile skips them (`nanmedian` / # `nanquantile`). mask3 = valid_mask.unsqueeze(-1) safe_xyz = torch.where(mask3, xyz, torch.full_like(xyz, float("nan"))) center = torch.nanmedian(safe_xyz, dim=1).values # (B, 3) d = torch.linalg.norm(xyz - center.unsqueeze(1), dim=-1) # (B, N) d = torch.where(valid_mask, d, torch.full_like(d, float("nan"))) scale = torch.nanquantile(d, 0.95, dim=1) # (B,) scale = scale.clamp(min=1e-3) return center.to(xyz.dtype), scale.to(xyz.dtype) def _robust_norm_params_all_valid_batched( xyz: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Fast median centre + P95 scale when every slot is valid.""" center = xyz.median(dim=1).values d = torch.linalg.norm(xyz - center.unsqueeze(1), dim=-1) scale = torch.quantile(d, 0.95, dim=1).clamp(min=1e-3) return center.to(xyz.dtype), scale.to(xyz.dtype) @torch.no_grad() def fps_subsample_stage2_batch( batch: dict, n_col: int, n_dep: int, seed_per_sample: Optional[torch.Tensor] = None, renormalize_bbox: bool = True, max_exact_iters_per_provenance: Optional[int] = None, ) -> dict: """Apply per-provenance FPS to a stage-2 batch, in place where possible. The batch is expected to come from ``stage2_collate_fn`` with the FPS variant: per-point tensors have shape ``(B, N_max, …)`` and a ``scene_valid_mask`` ``(B, N_max)`` marks real points vs. padding. Coordinates are in WORLD space (no bbox normalisation applied) — i.e. produced by ``stage2_row_to_sample(..., pre_subsample=False)``. Args: batch: dict with at least ``scene_xyz``, ``scene_type_ids``, ``scene_valid_mask``, ``bbox_center`` (filler placeholder, will be overwritten), ``bbox_scale``, ``bbox_R``. Optional per-point keys (gestalt/ADE/conf/RGB) are gathered if present. n_col / n_dep: per-provenance budgets. seed_per_sample: optional ``(B,)`` long for reproducible seeding. renormalize_bbox: if True, recompute ``bbox_center`` / ``bbox_scale`` from the FPS-subsampled cloud (matches the random path which also recomputes per augmented sample). If False, leaves the existing ``bbox_center`` / ``bbox_scale`` alone. max_exact_iters_per_provenance: optional cap for exact FPS anchors in each provenance. ``None`` preserves exact FPS. Smaller values enable the fast hybrid sampler. Returns: A new dict with the per-point tensors shrunk to ``(B, n_col + n_dep, …)`` and ``scene_xyz`` normalized into the post-subsample bbox. Non-point tensors (init_verts, verts_gt, etc.) are passed through un-changed. ``scene_valid_mask`` is set to all-True at the new size. """ xyz_world = batch["scene_xyz"] # (B, N, 3) WORLD type_ids = batch["scene_type_ids"] # (B, N) valid_mask = batch["scene_valid_mask"] # (B, N) B, N, _ = xyz_world.shape device = xyz_world.device idx = fps_per_provenance_indices( xyz_world, type_ids, valid_mask, n_col=n_col, n_dep=n_dep, seed_per_sample=seed_per_sample, max_exact_iters_per_provenance=max_exact_iters_per_provenance, ) # (B, K) K = idx.shape[1] arange_b = torch.arange(B, device=device).unsqueeze(1).expand(B, K) out = dict(batch) # Gather all per-point tensors. Index supports any trailing shape. for key in _STAGE2_SCENE_PER_POINT_KEYS: t = batch.get(key) if t is None: continue if t.dim() == 2: out[key] = t[arange_b, idx] else: # (B, N, C) — gather along dim=1 out[key] = t[arange_b, idx] # After FPS the valid mask is trivially all-True (we picked real or padded # slots; padded slots only sneak in when n_col / n_dep exceeds the # available count, in which case duplicate real picks are fine). out["scene_valid_mask"] = torch.ones((B, K), dtype=torch.bool, device=device) if renormalize_bbox: sub_xyz_world = out["scene_xyz"] # (B, K, 3) WORLD center, scale = _robust_norm_params_all_valid_batched(sub_xyz_world) out["scene_xyz"] = ((sub_xyz_world - center.unsqueeze(1)) / scale.view(B, 1, 1)).to(sub_xyz_world.dtype) out["bbox_center"] = center out["bbox_scale"] = scale # bbox_R is identity for the FPS path (no augmentation rotation # baked into the cache; matches the v1 random path during augment=False). out["bbox_R"] = torch.eye(3, device=device, dtype=sub_xyz_world.dtype) \ .unsqueeze(0).expand(B, 3, 3).contiguous() # Re-normalize init_verts / verts_gt into the new bbox so the model's # init and supervision land in the same frame as scene_xyz. init_world = batch.get("init_verts_world") if isinstance(init_world, torch.Tensor): out["init_verts"] = ( (init_world - center.unsqueeze(1)) / scale.view(B, 1, 1) ).to(init_world.dtype) verts_gt_world = batch.get("verts_gt_world") if isinstance(verts_gt_world, torch.Tensor): # verts_gt is (B, K_v, 4) — xyz + validity flag. xyz_g = verts_gt_world[..., :3] flag = verts_gt_world[..., 3:] xyz_g_n = (xyz_g - center.unsqueeze(1)) / scale.view(B, 1, 1) out["verts_gt"] = torch.cat([xyz_g_n.to(verts_gt_world.dtype), flag], dim=-1) return out