torch-pointcloud 0.0.2__tar.gz → 0.0.4__tar.gz

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/.gitignore

@@ -164,6 +164,7 @@ cython_debug/
 weights/**
 data/**
 docs/api/**
+docs/assets/data/
 docs/examples/**
 docs/scripts/**
 !docs/examples/*.ipynb

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: torch-pointcloud
-Version: 0.0.2
+Version: 0.0.4
 Summary: PyTorch Point Cloud models, scripts, pretrained weights and more
 Project-URL: Homepage, https://github.com/arthurdjn/pytorch-pointcloud
 Project-URL: Repository, https://github.com/arthurdjn/pytorch-pointcloud
@@ -51,7 +51,7 @@ Description-Content-Type: text/markdown
 [![tests](https://img.shields.io/github/actions/workflow/status/arthurdjn/pytorch-pointcloud/test.yml?branch=main&label=tests&labelColor=11001C&logo=github&logoColor=white)](https://github.com/arthurdjn/pytorch-pointcloud/actions/workflows/test.yml)
 [![pypi](https://img.shields.io/pypi/v/torch-pointcloud?color=EE4C2C&labelColor=11001C&logo=pypi&logoColor=white)](https://pypi.org/project/torch-pointcloud/)
 [![license](https://img.shields.io/badge/license-Apache--2.0-red.svg?color=EE4C2C&labelColor=11001C)](https://github.com/arthurdjn/pytorch-pointcloud/blob/main/LICENSE)
-[![doi](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.22159633-red.svg?color=EE4C2C&labelColor=11001C)](https://doi.org/10.5281/zenodo.22159633)
+[![doi](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.22159632-red.svg?color=EE4C2C&labelColor=11001C)](https://doi.org/10.5281/zenodo.22159632)
 [![docs](https://img.shields.io/badge/docs-pytorch--pointcloud.org-red.svg?color=EE4C2C&labelColor=11001C&logo=materialformkdocs&logoColor=white)](https://pytorch-pointcloud.org/)  
 
 </div>
@@ -158,7 +158,7 @@ Every checkpoint ships the transform that turns a raw point cloud into what the
 # Requires torch-scatter, torch-cluster, spconv
 model, info = tp.create_model("ptv3-base.scannet20.pointcept", task="segmentation", pretrained=True, return_info=True)
 info["transform"]  # the preprocessing pipeline of that checkpoint
-info["weights"]["metrics"]  # {"mIoU": 76.29}
+info["weights"]["metrics"]  # {"mIoU": 77.40}
 
 tp.list_models("pointnext*")  # every registered PointNeXt config
 tp.list_models(task="detection", pretrained=True)  # all detection checkpoints
@@ -186,7 +186,7 @@ If you find this project useful, please consider citing:
   title = {PyTorch PointCloud},
   year = {2026},
   url = {https://github.com/arthurdjn/pytorch-pointcloud},
-  doi = {10.5281/zenodo.22159633},
+  doi = {10.5281/zenodo.22159632},
   license = {Apache-2.0}
 }
 ```

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/README.md

@@ -8,7 +8,7 @@
 [![tests](https://img.shields.io/github/actions/workflow/status/arthurdjn/pytorch-pointcloud/test.yml?branch=main&label=tests&labelColor=11001C&logo=github&logoColor=white)](https://github.com/arthurdjn/pytorch-pointcloud/actions/workflows/test.yml)
 [![pypi](https://img.shields.io/pypi/v/torch-pointcloud?color=EE4C2C&labelColor=11001C&logo=pypi&logoColor=white)](https://pypi.org/project/torch-pointcloud/)
 [![license](https://img.shields.io/badge/license-Apache--2.0-red.svg?color=EE4C2C&labelColor=11001C)](https://github.com/arthurdjn/pytorch-pointcloud/blob/main/LICENSE)
-[![doi](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.22159633-red.svg?color=EE4C2C&labelColor=11001C)](https://doi.org/10.5281/zenodo.22159633)
+[![doi](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.22159632-red.svg?color=EE4C2C&labelColor=11001C)](https://doi.org/10.5281/zenodo.22159632)
 [![docs](https://img.shields.io/badge/docs-pytorch--pointcloud.org-red.svg?color=EE4C2C&labelColor=11001C&logo=materialformkdocs&logoColor=white)](https://pytorch-pointcloud.org/)  
 
 </div>
@@ -115,7 +115,7 @@ Every checkpoint ships the transform that turns a raw point cloud into what the
 # Requires torch-scatter, torch-cluster, spconv
 model, info = tp.create_model("ptv3-base.scannet20.pointcept", task="segmentation", pretrained=True, return_info=True)
 info["transform"]  # the preprocessing pipeline of that checkpoint
-info["weights"]["metrics"]  # {"mIoU": 76.29}
+info["weights"]["metrics"]  # {"mIoU": 77.40}
 
 tp.list_models("pointnext*")  # every registered PointNeXt config
 tp.list_models(task="detection", pretrained=True)  # all detection checkpoints
@@ -143,7 +143,7 @@ If you find this project useful, please consider citing:
   title = {PyTorch PointCloud},
   year = {2026},
   url = {https://github.com/arthurdjn/pytorch-pointcloud},
-  doi = {10.5281/zenodo.22159633},
+  doi = {10.5281/zenodo.22159632},
   license = {Apache-2.0}
 }
 ```

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/pyproject.toml

@@ -16,7 +16,7 @@ include = [
 
 [project]
 name = "torch-pointcloud"
-version = "0.0.2"
+version = "0.0.4"
 description = "PyTorch Point Cloud models, scripts, pretrained weights and more"
 readme = "README.md"
 license = "Apache-2.0"

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/src/torch_pointcloud/inferers/_utils.py

@@ -1,6 +1,6 @@
-"""Shared inferer helpers: dict indexing, chunk splitting, and Gaussian distance weighting."""
+"""Shared inferer helpers: dict indexing, per-fragment transforms, chunk splitting, and Gaussian distance weighting."""
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -21,6 +21,47 @@ def index_select_dict(data: Dict[str, Any], idx: Tensor, n_points: int) -> Dict[
     return out
 
 
+def apply_transform(
+    sample: Dict[str, Any],
+    transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]],
+    *,
+    pos_key: str,
+    inverse_key: Optional[str],
+) -> Tuple[Dict[str, Any], Optional[Tensor]]:
+    r"""Run the per-fragment `transform` over `sample` and return it with its source-to-predictor row map.
+
+    Any value already at `inverse_key` is dropped before `transform` runs, so a scene-level `inverse` never
+    becomes the prior the fragment's own map composes through. The map the transform writes under `inverse_key`
+    is popped from the returned sample, so it never reaches the predictor.
+
+    Args:
+        sample: Data dict of one fragment (block, sub-cloud, sphere).
+        transform: Callable applied to `sample`, or `None` to leave it unchanged.
+        pos_key: Dict key for the position tensor, whose row count defines the fragment size.
+        inverse_key: Dict key under which a row-altering `transform` records its long index map of shape
+            $(N_\text{source},)$ with values in $[0, N_\text{predictor})$. `None` when rows are preserved.
+
+    Returns:
+        The transformed sample and its inverse map, `None` when the transform wrote none. Predictions of shape
+        $(N_\text{predictor}, C)$ index with the map back to the $N_\text{source}$ rows of `sample`.
+
+    Raises:
+        ValueError: If `transform` changes the row count without writing a map under `inverse_key`.
+    """
+    n_source = int(sample[pos_key].size(0))
+    sample = {key: value for key, value in sample.items() if key != inverse_key}
+    if transform is not None:
+        sample = transform(sample)
+    inverse_map = sample.pop(inverse_key, None) if inverse_key is not None else None
+    n_predictor = int(sample[pos_key].size(0))
+    if inverse_map is None and n_predictor != n_source:
+        raise ValueError(
+            f"`transform` changed the fragment's row count ({n_source} -> {n_predictor}) without recording an "
+            f"index map. Make it write one (e.g. `dst_inverse_key`) and pass that key as `inverse_key`."
+        )
+    return sample, inverse_map
+
+
 def check_batch_alignment(pos: Tensor, batch: Tensor, pos_key: str, batch_key: str) -> None:
     """Raise when the per-point batch index does not line up with the positions row for row.
 
@@ -35,7 +76,7 @@ def check_batch_alignment(pos: Tensor, batch: Tensor, pos_key: str, batch_key: s
         )
 
 
-def split_chunks(n: int, max_size: Optional[int], rng: torch.Generator) -> List[Tensor]:
+def split_chunks(n: int, max_size: Optional[int], rng: Optional[torch.Generator], device: torch.device) -> List[Tensor]:
     r"""Partition `range(n)` into index chunks of at most `max_size` points each.
 
     When `max_size` is `None` or `n` is within `max_size`, returns a single in-order
@@ -46,13 +87,14 @@ def split_chunks(n: int, max_size: Optional[int], rng: torch.Generator) -> List[
     Args:
         n: Number of indices to partition.
         max_size: Maximum points per chunk, or `None` for a single unsplit chunk.
-        rng: Generator for the permutation. Only consumed when a split is needed.
+        rng: Generator for the permutation, or `None` for the global generator. Only consumed when a split
+            is needed.
+        device: Device of the returned indices, which a given `rng` must live on.
 
     Returns:
-        List of 1-D `long` index tensors on `rng`'s device. The chunks partition
+        List of 1-D `long` index tensors on `device`. The chunks partition
         `range(n)` exactly: every index appears in one and only one chunk.
     """
-    device = rng.device
     if max_size is None:
         return [torch.arange(n, device=device)]
     if n > max_size:

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/src/torch_pointcloud/inferers/inferer.py

@@ -27,21 +27,27 @@ class Inferer(metaclass=ABCMeta):
 
     `data` is a packed-batch dict (at minimum containing position and batch indices).
     `predictor` is any callable taking such a dict and returning per-point logits of shape
-    $(N, C_\text{out})$. Inferers are stateless with respect to the scene, so
-    one instance can be reused across scenes and wrapped by another inferer.
-
-    Every concrete inferer exposes a knob controlling whether partial predictions are
-    converted to softmax probabilities before aggregation. The defaults differ:
-
-    | Inferer                  | Parameter     | Default | Aggregated quantity                                          |
-    | ------------------------ | ------------- | ------- | ------------------------------------------------------------ |
-    | `SimpleInferer`          | `softmax`     | `False` | predictor output as-is                                       |
-    | `SlidingWindowInferer`   | `softmax`     | `True`  | softmax probabilities per block (`"max"` / `"vote"` always)  |
-    | `KNNWindowInferer`       | `softmax`     | `False` | raw logits (`"weighted_mean"`); always probabilities (`"ema"`) |
-    | `VoxelPartitionInferer`  | `softmax`     | `False` | raw logits per pass                                          |
-    | `PotentialSphereInferer` | none          |         | always an EMA of softmax probabilities                        |
-    | `TTAInferer`             | `ema_softmax` | `True`  | base output as-is (`"mean"`); softmax of base output (`"ema"`) |
-    | `PartRefinementInferer`  | none          |         | one-hot refined labels of the base output's argmax           |
+    $(N, C_\text{out})$. Inferers keep nothing of the scene, so one instance can be
+    reused across scenes and wrapped by another inferer.
+
+    Inferers that draw random numbers take a `seed`. `None` follows the global generator,
+    like the transforms do, so one `torch.manual_seed` (or `seed_everything`) seeds the
+    whole evaluation. An int gives the inferer its own stream, offset by the number of
+    calls the instance has made: successive calls (e.g. `TTAInferer` votes) draw different
+    numbers, and a fresh instance replays the same sequence.
+
+    The `softmax` parameter controls whether partial predictions are converted to softmax
+    probabilities before aggregation, and `aggregate` how they are combined:
+
+    | Inferer                  | `softmax` | Aggregated quantity                                         |
+    | ------------------------ | --------- | ----------------------------------------------------------- |
+    | `SimpleInferer`          | `False`   | predictor output as-is                                      |
+    | `SlidingWindowInferer`   | `True`    | softmax probabilities per block (`"max"` / `"vote"` always) |
+    | `KNNWindowInferer`       | `False`   | raw logits (`"mean"`); always probabilities (`"ema"`)       |
+    | `VoxelPartitionInferer`  | `False`   | raw logits per pass                                         |
+    | `PotentialSphereInferer` | none      | always an EMA of softmax probabilities                      |
+    | `TTAInferer`             | `False`   | base output as-is                                           |
+    | `PartRefinementInferer`  | none      | one-hot refined labels of the base output's argmax          |
 
     When the input scene is empty ($N = 0$), inferers that never call the predictor
     return a $(0, 0)$ tensor (the channel count cannot be inferred without a predictor

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/src/torch_pointcloud/inferers/knn_window.py

@@ -19,11 +19,11 @@ from tqdm import tqdm
 
 from torch_pointcloud.utils.data import DataKeys
 
-from ._utils import check_batch_alignment, gaussian_weights, index_select_dict
+from ._utils import apply_transform, check_batch_alignment, gaussian_weights, index_select_dict
 from .inferer import Inferer
 
 WindowMode = Literal["constant", "gaussian"]
-AggregateMode = Literal["weighted_mean", "ema"]
+AggregateMode = Literal["mean", "ema"]
 
 
 def _knn_centers(pos_src: Tensor, centers: Tensor, k: int) -> Tensor:
@@ -62,12 +62,13 @@ def knn_window_inference(
     overlap: float = 0.5,
     mode: WindowMode = "constant",
     sigma_scale: float = 0.125,
-    aggregate: AggregateMode = "weighted_mean",
+    aggregate: AggregateMode = "mean",
     ema_smoothing: float = 0.95,
     softmax: bool = False,
     transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
     pos_key: str = DataKeys.POS,
     batch_key: str = DataKeys.BATCH,
+    inverse_key: Optional[str] = None,
     progress: bool = False,
     seed: Optional[int] = None,
 ) -> Tensor:
@@ -81,7 +82,7 @@ def knn_window_inference(
 
     `aggregate` controls how overlapping window predictions are combined:
 
-    - `"weighted_mean"`: accumulates distance-weighted logits and divides by total
+    - `"mean"`: accumulates distance-weighted logits and divides by total
       weight at the end, producing a weighted-average logit tensor.
     - `"ema"`: per-update softmax EMA
       ($\text{new} = \alpha \cdot \text{old} + (1 - \alpha) \cdot \text{softmax}(\text{logits})$),
@@ -109,10 +110,10 @@ def knn_window_inference(
         mode: Distance weighting for each window. `"constant"` gives equal weight
             to every point; `"gaussian"` weights by $\exp(-d^2 / 2\sigma^2)$ with
             $\sigma = \text{sigma\_scale} \cdot \max_i d_i$. `"gaussian"` requires
-            `aggregate="weighted_mean"`; EMA updates ignore distance weights.
+            `aggregate="mean"`; EMA updates ignore distance weights.
         sigma_scale: Gaussian sigma scale factor (only used when `mode="gaussian"`).
         aggregate: How predictions from overlapping windows are combined.
-            `"weighted_mean"`: weighted-average logits (divide by total weight at end).
+            `"mean"`: weighted-average logits (divide by total weight at end).
             `"ema"`: softmax EMA
             ($\text{new} = \alpha \cdot \text{old} + (1 - \alpha) \cdot \text{softmax}(\text{logits})$);
             use `sw_batch_size=1` with EMA to match the reference evaluation protocol.
@@ -123,14 +124,23 @@ def knn_window_inference(
             softmax probabilities.
         transform: Optional callable applied to each window's data dict before the
             predictor (typical example: `T.Shift(keys=DataKeys.POS, method="centroid")`).
+            If it changes the row count (pad, voxelize, ...) it must record a
+            source-to-predictor index map under `inverse_key` so the inferer can gather
+            predictions back to the window's points.
         pos_key: Dict key for the position tensor.
         batch_key: Dict key for the per-point batch index.
+        inverse_key: Dict key under which a row-altering `transform` records a
+            source-to-predictor long index map of shape $(k,)$ with values in
+            $[0, N_\text{predictor})$. Any scene-level value at this key is dropped before
+            `transform` runs, and the map is popped before the predictor is called. Leave
+            `None` when the transform preserves row count.
         progress: If `True`, show a `tqdm` progress bar per batch element.
-        seed: RNG seed for the per-point initial possibility scores.
+        seed: RNG seed for the per-point initial possibility scores. `None` draws from the global generator, so
+            `torch.manual_seed` seeds the inferer together with the transforms.
 
     Returns:
         Per-point output tensor of shape $(N, C_\text{out})$. Aggregation produces
-        logits when `aggregate="weighted_mean"` (probabilities with `softmax=True`)
+        logits when `aggregate="mean"` (probabilities with `softmax=True`)
         and probabilities when `aggregate="ema"`. An empty scene ($N = 0$) returns
         a $(0, 0)$ tensor: the predictor is never called, so the channel count
         cannot be inferred.
@@ -143,8 +153,8 @@ def knn_window_inference(
         raise ValueError(f"`overlap` must be in (0, 1), got {overlap}.")
     if mode not in ("constant", "gaussian"):
         raise ValueError(f"`mode` must be 'constant' or 'gaussian', got {mode!r}.")
-    if aggregate not in ("weighted_mean", "ema"):
-        raise ValueError(f"`aggregate` must be 'weighted_mean' or 'ema', got {aggregate!r}.")
+    if aggregate not in ("mean", "ema"):
+        raise ValueError(f"`aggregate` must be 'mean' or 'ema', got {aggregate!r}.")
     if sw_batch_size < 1:
         raise ValueError(f"`sw_batch_size` must be >= 1, got {sw_batch_size}.")
     if not 0.0 <= ema_smoothing < 1.0:
@@ -152,7 +162,7 @@ def knn_window_inference(
     if aggregate == "ema" and mode == "gaussian":
         raise ValueError(
             "`mode='gaussian'` is incompatible with `aggregate='ema'`: EMA updates blend by `ema_smoothing`, not "
-            "by per-point distance weights. Use `aggregate='weighted_mean'` or `mode='constant'`."
+            "by per-point distance weights. Use `aggregate='mean'` or `mode='constant'`."
         )
     if aggregate == "ema" and sw_batch_size > 1:
         warnings.warn(
@@ -170,9 +180,7 @@ def knn_window_inference(
     n_total = pos.size(0)
     output: Optional[Tensor] = None
 
-    rng = torch.Generator(device=device)
-    if seed is not None:
-        rng.manual_seed(int(seed))
+    rng = None if seed is None else torch.Generator(device=device).manual_seed(int(seed))
 
     for b in torch.unique(batch).tolist():
         idx_b = torch.where(batch == b)[0]
@@ -201,37 +209,39 @@ def knn_window_inference(
                 local_idxs = _knn_centers(pos_b, center_pos, k)
                 flat_idxs = local_idxs.reshape(-1)
                 per_window_dicts: List[Dict[str, Any]] = []
+                inverse_maps: List[Optional[Tensor]] = []
                 for w_i in range(sw):
                     w_idx = local_idxs[w_i]
                     wd = index_select_dict(data_b, w_idx, n_b)
                     wd[batch_key] = torch.full((k,), w_i, device=device, dtype=torch.long)
-                    if transform is not None:
-                        wd = transform(wd)
-                        if int(wd[pos_key].size(0)) != k:
-                            raise ValueError(
-                                f"`transform` must preserve each window's row count; got {k} -> "
-                                f"{int(wd[pos_key].size(0))} rows."
-                            )
-
+                    wd, inverse_map = apply_transform(wd, transform, pos_key=pos_key, inverse_key=inverse_key)
+                    wd[batch_key] = torch.full((int(wd[pos_key].size(0)),), w_i, device=device, dtype=torch.long)
                     per_window_dicts.append(wd)
+                    inverse_maps.append(inverse_map)
 
+                sizes = [int(wd[pos_key].size(0)) for wd in per_window_dicts]
                 window_data: Dict[str, Any] = {}
                 for key in per_window_dicts[0]:
                     values = [wd[key] for wd in per_window_dicts]
-                    if torch.is_tensor(values[0]) and values[0].dim() > 0 and values[0].size(0) == k:
+                    if torch.is_tensor(values[0]) and values[0].dim() > 0 and values[0].size(0) == sizes[0]:
                         window_data[key] = torch.cat(values, dim=0)
                     else:
                         window_data[key] = values[0]
 
                 window_logits = predictor(window_data).to(device)
                 num_classes = int(window_logits.size(-1))
-                window_logits = window_logits.reshape(sw, k, num_classes)
+                window_logits = torch.stack(
+                    [
+                        part if inverse_map is None else part[inverse_map]
+                        for part, inverse_map in zip(window_logits.split(sizes), inverse_maps)
+                    ]
+                )  # (sw, k, C)
 
                 if output is None:
                     output = torch.zeros(n_total, num_classes, device=device, dtype=torch.float32)
                 if scores_b is None:
                     scores_b = torch.zeros(n_b, num_classes, device=device, dtype=torch.float32)
-                    if aggregate == "weighted_mean":
+                    if aggregate == "mean":
                         weights_b = torch.zeros(n_b, device=device, dtype=torch.float32)
 
                 distances = torch.linalg.norm(pos_b[local_idxs] - center_pos.unsqueeze(1), dim=-1)
@@ -289,7 +299,9 @@ class KNNWindowInferer(Inferer):
     the least-covered points until all points are covered. Reuse the same instance
     across scenes; compose with `TTAInferer` for multi-augmentation averaging.
 
-    All parameters are forwarded verbatim to `knn_window_inference`.
+    All parameters are forwarded verbatim to `knn_window_inference`, except `seed`, which is offset by
+    the number of calls the instance has made: repeated calls (e.g. `TTAInferer` votes) draw different random
+    numbers, and a fresh instance replays the same sequence.
 
     Example:
         ```python
@@ -308,12 +320,13 @@ class KNNWindowInferer(Inferer):
         overlap: float = 0.5,
         mode: WindowMode = "constant",
         sigma_scale: float = 0.125,
-        aggregate: AggregateMode = "weighted_mean",
+        aggregate: AggregateMode = "mean",
         ema_smoothing: float = 0.95,
         softmax: bool = False,
         transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
         pos_key: str = DataKeys.POS,
         batch_key: str = DataKeys.BATCH,
+        inverse_key: Optional[str] = None,
         progress: bool = False,
         seed: Optional[int] = None,
     ) -> None:
@@ -328,14 +341,18 @@ class KNNWindowInferer(Inferer):
         self.transform = transform
         self.pos_key = pos_key
         self.batch_key = batch_key
+        self.inverse_key = inverse_key
         self.progress = progress
         self.seed = seed
+        self._num_calls = 0
 
     def forward(
         self,
         data: Dict[str, Any],
         predictor: Callable[[Dict[str, Any]], Tensor],
     ) -> Tensor:
+        seed = None if self.seed is None else self.seed + self._num_calls
+        self._num_calls += 1
         return knn_window_inference(
             data,
             predictor=predictor,
@@ -350,6 +367,7 @@ class KNNWindowInferer(Inferer):
             transform=self.transform,
             pos_key=self.pos_key,
             batch_key=self.batch_key,
+            inverse_key=self.inverse_key,
             progress=self.progress,
-            seed=self.seed,
+            seed=seed,
         )

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.4}/src/torch_pointcloud/inferers/potential_sphere.py

@@ -15,7 +15,7 @@ from tqdm import tqdm
 from torch_pointcloud.utils.data import DataKeys, collate
 from torch_pointcloud.utils.ops import voxel_grid_fnv
 
-from ._utils import check_batch_alignment, index_select_dict
+from ._utils import apply_transform, check_batch_alignment, index_select_dict
 from .inferer import Inferer
 
 
@@ -23,7 +23,7 @@ def _next_sphere(
     pos_b: Tensor,
     coarse_pos: Tensor,
     potentials: Tensor,
-    rng: torch.Generator,
+    rng: Optional[torch.Generator],
     radius: float,
     jitter: float,
 ) -> Tuple[Tensor, Tensor]:
@@ -55,6 +55,7 @@ def potential_sphere_inference(
     transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
     pos_key: str = DataKeys.POS,
     batch_key: str = DataKeys.BATCH,
+    inverse_key: Optional[str] = None,
     progress: bool = False,
     seed: Optional[int] = None,
 ) -> Tensor:
@@ -91,12 +92,19 @@ def potential_sphere_inference(
         sw_batch_size: Number of spheres packed into one predictor call. Centers are still drawn one at a
             time with the potentials updated in between, as the reference sampler does.
         transform: Optional per-sphere callable applied to the centered sphere dict before the predictor.
-            The transform must preserve the sphere's row count and keep positions centered on the sphere (the
-            `inner_ratio` mask is evaluated on the transformed positions, as the reference does).
+            The transform must keep positions centered on the sphere (the `inner_ratio` mask is evaluated on
+            the transformed positions, as the reference does). If it changes the row count (pad, voxelize, ...)
+            it must record a source-to-predictor index map under `inverse_key` so the inferer can gather
+            predictions and the mask back to the sphere's points.
         pos_key: Dict key for the position tensor.
         batch_key: Dict key for the per-point batch index.
+        inverse_key: Dict key under which a row-altering `transform` records a source-to-predictor long index
+            map of shape $(N_\text{sphere},)$ with values in $[0, N_\text{predictor})$. Any scene-level value at
+            this key is dropped before `transform` runs, and the map is popped before the predictor is called.
+            Leave `None` when the transform preserves row count.
         progress: If `True`, show a `tqdm` progress bar per batch element.
-        seed: Optional RNG seed for the initial potentials and the center jitter.
+        seed: RNG seed for the initial potentials and the center jitter. `None` draws from the global generator, so
+            `torch.manual_seed` seeds the inferer together with the transforms.
 
     Returns:
         Per-point score tensor of shape $(N, C)$: the EMA of softmax probabilities over the spheres covering
@@ -132,9 +140,7 @@ def potential_sphere_inference(
     n = pos.size(0)
     inner_sq = (inner_ratio * radius) ** 2
 
-    rng = torch.Generator(device=device)
-    if seed is not None:
-        rng.manual_seed(int(seed))
+    rng = None if seed is None else torch.Generator(device=device).manual_seed(int(seed))
 
     output: Optional[Tensor] = None
     for b in torch.unique(batch).tolist():
@@ -152,6 +158,7 @@ def potential_sphere_inference(
             while float(potentials.min().item()) < num_votes:
                 spheres: List[Dict[str, Any]] = []
                 sphere_idx: List[Tensor] = []
+                inverse_maps: List[Optional[Tensor]] = []
                 for _ in range(sw_batch_size):
                     if float(potentials.min().item()) >= num_votes:
                         break
@@ -162,16 +169,11 @@ def potential_sphere_inference(
 
                     sphere = index_select_dict(data_b, idx, n_b)
                     sphere[pos_key] = sphere[pos_key] - center
-                    if transform is not None:
-                        sphere = transform(sphere)
-                        if int(sphere[pos_key].size(0)) != int(idx.numel()):
-                            raise ValueError(
-                                f"`transform` must preserve each sphere's row count; got {int(idx.numel())} -> "
-                                f"{int(sphere[pos_key].size(0))} rows."
-                            )
+                    sphere, inverse_map = apply_transform(sphere, transform, pos_key=pos_key, inverse_key=inverse_key)
 
                     spheres.append(sphere)
                     sphere_idx.append(idx)
+                    inverse_maps.append(inverse_map)
 
                 if not spheres:
                     continue
@@ -182,11 +184,14 @@ def potential_sphere_inference(
                     scores_b = torch.zeros(n_b, int(probs.size(-1)), device=device, dtype=probs.dtype)
 
                 offset = 0
-                for sphere, idx in zip(spheres, sphere_idx):
-                    count = int(idx.numel())
+                for sphere, idx, inverse_map in zip(spheres, sphere_idx, inverse_maps):
+                    count = int(sphere[pos_key].size(0))
                     sphere_probs = probs[offset : offset + count]
                     offset += count
                     inner = sphere[pos_key].square().sum(dim=-1) < inner_sq
+                    if inverse_map is not None:
+                        sphere_probs = sphere_probs[inverse_map]
+                        inner = inner[inverse_map]
                     rows = idx[inner]
                     scores_b[rows] = ema_smoothing * scores_b[rows] + (1.0 - ema_smoothing) * (
                         sphere_probs[inner].to(scores_b.dtype)
@@ -225,7 +230,9 @@ class PotentialSphereInferer(Inferer):
     each sphere's softmax predictions into the running per-point scores by an exponential moving average,
     until every region has been covered about `num_votes` times.
 
-    All parameters are forwarded verbatim to `potential_sphere_inference`.
+    All parameters are forwarded verbatim to `potential_sphere_inference`, except `seed`, which is offset by
+    the number of calls the instance has made: repeated calls (e.g. `TTAInferer` votes) draw different random
+    numbers, and a fresh instance replays the same sequence.
 
     Example:
         ```{.python notest}
@@ -248,6 +255,7 @@ class PotentialSphereInferer(Inferer):
         transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
         pos_key: str = DataKeys.POS,
         batch_key: str = DataKeys.BATCH,
+        inverse_key: Optional[str] = None,
         progress: bool = False,
         seed: Optional[int] = None,
     ) -> None:
@@ -261,14 +269,18 @@ class PotentialSphereInferer(Inferer):
         self.transform = transform
         self.pos_key = pos_key
         self.batch_key = batch_key
+        self.inverse_key = inverse_key
         self.progress = progress
         self.seed = seed
+        self._num_calls = 0
 
     def forward(
         self,
         data: Dict[str, Any],
         predictor: Callable[[Dict[str, Any]], Tensor],
     ) -> Tensor:
+        seed = None if self.seed is None else self.seed + self._num_calls
+        self._num_calls += 1
         return potential_sphere_inference(
             data,
             predictor=predictor,
@@ -282,6 +294,7 @@ class PotentialSphereInferer(Inferer):
             transform=self.transform,
             pos_key=self.pos_key,
             batch_key=self.batch_key,
+            inverse_key=self.inverse_key,
             progress=self.progress,
-            seed=self.seed,
+            seed=seed,
         )