Overview

Iterative KNN-window inference for large-scale point cloud segmentation.

Maintains a per-point coverage score initialized to small random noise. Each iteration selects the sw_batch_size least-covered points as window centers, crops their \(k\) nearest neighbors, runs predictor on the packed crop, and accumulates per-point predictions weighted by distance to the center. The loop ends once every point's coverage score exceeds overlap.

aggregate controls how overlapping window predictions are combined:

• "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})\)), outputting calibrated probabilities without a final softmax step.

Parameters:

• data (Dict[str, Any]) –

  Dict of per-point tensors. Must contain pos (shape \((N, D)\)) and batch (shape \((N,)\)). Any additional tensor whose first dim equals \(N\) (e.g. color, intensity, segment) is automatically sliced to the active window. Non-tensor entries and tensors with a different leading dim (scalar metadata) flow through unchanged.

• predictor (Callable[[Dict[str, Any]], Tensor]) –

  Callable taking a per-window data dict and returning per-point logits of shape \((M, C_\text{out})\), where \(M\) is the (packed) total point count of the windowed batch and the per-window batch index lives at window[batch_key]. The output channel count \(C_\text{out}\) is inferred from the first call.

• roi_num_points (int, default: 65536 ) –

  Number of points per window (the window size \(k\)).

• sw_batch_size (int, default: 1 ) –

  Number of windows packed into one predictor call. Higher values reduce launch overhead at the cost of more memory.

• overlap (float, default: 0.5 ) –

  Coverage threshold in \((0, 1)\). The loop stops once every point has accumulated at least overlap possibility mass. Higher values give more thorough coverage at the cost of more iterations. 0.0 is rejected: the initial coverage scores already satisfy a zero threshold, so no window would ever be predicted.

• mode (WindowMode, default: 'constant' ) –

  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="mean"; EMA updates ignore distance weights.

• sigma_scale (float, default: 0.125 ) –

  Gaussian sigma scale factor (only used when mode="gaussian").

• aggregate (AggregateMode, default: 'mean' ) –

  How predictions from overlapping windows are combined. "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.

• ema_smoothing (float, default: 0.95 ) –

  EMA factor \(\alpha \in [0, 1)\) used when aggregate="ema".

• softmax (bool, default: False ) –

  If True, softmax each window's logits before the weighted-mean accumulation, so overlapping windows average probabilities instead of raw logits. Ignored when aggregate="ema", which always accumulates softmax probabilities.

• transform (Optional[Callable[[Dict[str, Any]], Dict[str, Any]]], default: None ) –

  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 (str, default: POS ) –

  Dict key for the position tensor.

• batch_key (str, default: BATCH ) –

  Dict key for the per-point batch index.

• inverse_key (Optional[str], default: None ) –

  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 (bool, default: False ) –

  If True, show a tqdm progress bar per batch element.

• seed (Optional[int], default: None ) –

  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:

• Tensor –

  Per-point output tensor of shape \((N, C_\text{out})\). Aggregation produces

• Tensor –

  logits when aggregate="mean" (probabilities with softmax=True)

• Tensor –

  and probabilities when aggregate="ema". An empty scene (\(N = 0\)) returns

• Tensor –

  a \((0, 0)\) tensor: the predictor is never called, so the channel count

• Tensor –

  cannot be inferred.

Nearest-neighbor refinement of part labels on top of another inferer.

Runs base, takes the per-point argmax, and for every shape re-assigns the labels that are implausible: a predicted part with fewer than min_count points, or a part the shape's category does not own. Each such label is refined in turn (ascending label order): its points take the majority label of their num_neighbors nearest points of the same shape, the label under refinement excluded from the vote, with the already-refined labels feeding the next votes. This is the post-processing of the PointNeXt ShapeNetPart protocol.

Parameters:

• data (Dict[str, Any]) –

  Dict of per-point tensors. Must contain pos (shape \((N, D)\)), batch (shape \((N,)\)) and the per-shape category under category_key.

• predictor (Callable[[Dict[str, Any]], Tensor]) –

  Callable mapping a data dict to per-point part scores of shape \((N, C)\).

• base (Optional[Inferer], default: None ) –

  Inferer running predictor; defaults to SimpleInferer (one forward on the whole batch).

• part_ids (Optional[Sequence[Sequence[int]]], default: None ) –

  Part labels owned by each category; defaults to the 16-category / 50-part ShapeNetPart table (ShapeNetPart.seg_ids).

• min_count (int, default: 10 ) –

  Predicted parts with fewer points than this are refined.

• num_neighbors (int, default: 11 ) –

  Number of nearest neighbors (the point itself included) voting on the new label.

• category_key (str, default: CATEGORY ) –

  Dict key of the per-shape category, one-hot \((B, K)\) or index \((B,)\).

• pos_key (str, default: POS ) –

  Dict key for the position tensor.

• batch_key (str, default: BATCH ) –

  Dict key for the per-point batch index.

Returns:

• Tensor –

  One-hot refined labels of shape \((N, C)\), so a metric's argmax recovers the refined labels. An empty

• Tensor –

  scene (\(N = 0\)) returns the base output unchanged.

Potential-driven sphere voting for large-scale point cloud segmentation.

The scene is covered by spheres of radius \(r\) whose centers are chosen where the cloud has been seen the least: a coarse grid of potentials (one scalar per potential_size cell, initialized with a small random value) tracks coverage, each sphere is centered on the cell with the lowest potential (plus a Gaussian jitter) and raises the potentials of the cells it covers by the Tukey window \((1 - d^2 / r^2)^2\). The predictor runs on every sphere and its softmax probabilities are blended into the running per-point scores by an exponential moving average, restricted to the points within inner_ratio \(\cdot r\) of the center where the sphere's context is complete. The loop stops once every cell's potential reaches num_votes, so each region has been predicted about that many times.

This is the test protocol of KPConv: Flexible and Deformable Convolution for Point Clouds (Thomas et al., 2019) (radius-defined input spheres, test_smooth EMA, potential sampling), and it composes with any model that consumes a packed sphere: the per-sphere transform sees the centered sphere dict and can add the reference's stochastic test-time augmentation and the model's feature stack. Points that no sphere reaches keep all-zero scores.

Parameters:

• data (Dict[str, Any]) –

  Dict of per-point tensors. Must contain pos (shape \((N, D)\)) and batch (shape \((N,)\)); extra per-point tensors are sliced to the active sphere automatically.

• predictor (Callable[[Dict[str, Any]], Tensor]) –

  Callable mapping a packed sphere dict to per-point logits of shape \((M, C)\).

• radius (float) –

  Sphere radius, in the units of pos.

• num_votes (float, default: 10.0 ) –

  Potential threshold ending the loop, i.e. the number of times every region is covered (KPConv reports its numbers at the first multiple of \(10\)).

• potential_size (Optional[float], default: None ) –

  Cell size of the coarse potential grid. Defaults to radius / 10.

• jitter (Optional[float], default: None ) –

  Standard deviation of the Gaussian jitter added to each sphere center, clipped at radius / 2. Defaults to radius / 10; 0 disables it.

• inner_ratio (float, default: 0.7 ) –

  Fraction of radius inside which the sphere's predictions are kept.

• ema_smoothing (float, default: 0.95 ) –

  EMA factor \(\alpha \in [0, 1)\) of the score update \(\text{new} = \alpha \cdot \text{old} + (1 - \alpha) \cdot \text{softmax}(\text{logits})\).

• sw_batch_size (int, default: 1 ) –

  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[Callable[[Dict[str, Any]], Dict[str, Any]]], default: None ) –

  Optional per-sphere callable applied to the centered sphere dict before the predictor. 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 (str, default: POS ) –

  Dict key for the position tensor.

• batch_key (str, default: BATCH ) –

  Dict key for the per-point batch index.

• inverse_key (Optional[str], default: None ) –

  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 (bool, default: False ) –

  If True, show a tqdm progress bar per batch element.

• seed (Optional[int], default: None ) –

  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:

• Tensor –

  Per-point score tensor of shape \((N, C)\): the EMA of softmax probabilities over the spheres covering

• Tensor –

  each point; points no sphere reaches keep all-zero scores. An empty scene (\(N = 0\)) returns a

• Tensor –

  \((0, 0)\) tensor: the predictor is never called, so the channel count cannot be inferred.

Sliding-window inference for large-scale point cloud segmentation.

Places block centers on a regular grid with step \(\text{block\_size} \cdot (1 - \text{overlap})\) and calls predictor once per non-empty block. Each point's predictions from all covering blocks are combined by aggregate: a distance-weighted average ("mean"), the single most confident prediction ("max"), or a count of hard votes ("vote").

With overlap=0 and mode="constant", each point lands in exactly one block and the weight division is a no-op.

Parameters:

• data (Dict[str, Any]) –

  Dict of per-point tensors. Must contain pos (shape \((N, D)\)) and batch (shape \((N,)\)). Extra per-point tensors are sliced to the active block automatically. Non-tensor entries and tensors with a different leading dim flow through unchanged.

• predictor (Callable[[Dict[str, Any]], Tensor]) –

  Callable taking a per-block data dict and returning logits of shape \((M, C_\text{out})\), where \(M\) is the number of points in the block. The output channel count \(C_\text{out}\) is inferred from the first call.

• block_size (float) –

  Side length of each cubic block, in the same units as data[pos_key].

• overlap (float, default: 0.0 ) –

  Fraction of block_size shared between adjacent blocks, in \([0, 1)\). 0.0 gives a strict non-overlapping partition; 0.5 means adjacent block centers are spaced half a block apart.

• mode (WindowMode, default: 'constant' ) –

  Per-point weight within each block. "constant" gives equal weight to all points in the block. "gaussian" weights by \(\exp(-d^2 / 2\sigma^2)\) where \(d\) is the distance to the block center across the tiled axes and \(\sigma = \text{sigma\_scale} \cdot \text{block\_size} \cdot \sqrt{D_\text{tiled}} / 2\) with \(D_\text{tiled}\) the number of tiled axes.

• sigma_scale (float, default: 0.125 ) –

  Gaussian sigma scale factor. Only used when mode="gaussian".

• roi_num_points (Optional[int], default: None ) –

  Maximum points in one batch element. A block with more is split into random sub-batches, each predicted exactly once per block pass. None passes the whole block as one element. Use it when the architecture needs a bounded or fixed \(N\): pair it with a DivisiblePad-style transform that pads each block to a multiple of roi_num_points and writes its source-to-padded index map under inverse_key, and every element is then exactly roi_num_points long.

• sw_batch_size (int, default: 1 ) –

  Number of elements collated into one predictor call. Purely a throughput knob: it changes the number of calls, never the output and never the size of an element. Set together with roi_num_points, a call holds at most sw_batch_size elements of at most roi_num_points points each. Small blocks are dominated by per-call overhead, where a larger value is several times faster. Output-neutrality assumes predictor respects the packed batch index; one that pools globally without it will differ at sw_batch_size > 1. Peak memory holds sw_batch_size transformed windows at once, so it scales with block_size rather than with roi_num_points.

• softmax (bool, default: False ) –

  If True, softmax each block's logits before accumulating. Use True when averaging predictions across multiple blocks or TTA passes. Set False to accumulate raw logits. "max" and "vote" aggregation always read confidences off the softmax.

• aggregate (AggregateMode, default: 'mean' ) –

  How the predictions of the blocks covering a point are combined. "mean": distance-weighted average of the (softmax) predictions. "max": winner-takes-all, each point keeps the prediction of the block that is most confident about it (the PVCNN / PointCNN scene merge). "vote": each block casts one hard vote (its argmax) per point and the output holds the weighted vote fractions, so argmax is the majority label (the PointNet++ protocol).

• transform (Optional[Callable[[Dict[str, Any]], Dict[str, Any]]], default: None ) –

  Optional callable applied to each block's data dict before the predictor. The transform sees the whole block; 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 original block points.

• dims (Optional[Sequence[int]], default: None ) –

  Axes (indices into pos's last dim) to tile. None tiles every axis (cubic blocks). Pass (0, 1) for 2D tiling that leaves the third axis spanning the full scene height.

• padding (float, default: 0.0 ) –

  Extra margin (in pos units) extending each block's membership on every tiled axis. Useful for including a thin context guard band.

• pos_key (str, default: POS ) –

  Dict key for the position tensor.

• batch_key (str, default: BATCH ) –

  Dict key for the per-point batch index.

• block_bbox_key (str, default: BLOCK_BBOX ) –

  Dict key under which the block bounding box is exposed to the transform callable.

• inverse_key (Optional[str], default: None ) –

  Dict key under which a row-altering transform records a source-to-predictor long index map of shape \((N_\text{block},)\) with values in \([0, N_\text{window})\), where \(N_\text{block}\) is the pre-transform block size and \(N_\text{window}\) is the post-transform size. When set, any scene-level value at this key is dropped from the window before transform runs, so a registered pipeline's inverse never becomes the prior the block map composes through; the inferer then pops the block map before calling the predictor and gathers predictions back to block-local rows. Leave None when the transform preserves row count, or when no transform is used.

• progress (bool, default: False ) –

  If True, show a tqdm progress bar per batch element.

• seed (Optional[int], default: None ) –

  RNG seed for sub-batch permutations when roi_num_points is set. None draws from the global generator, so torch.manual_seed seeds the inferer together with the transforms.

Returns:

• Tensor –

  Per-point output tensor of shape \((N, C_\text{out})\): with

• Tensor –

  aggregate="mean" a distance-weighted average of softmax probabilities

• Tensor –

  when softmax=True or of raw logits when softmax=False; with

• Tensor –

  "max" the most confident block's probabilities; with "vote" the

• Tensor –

  per-class vote fractions. An empty scene (\(N = 0\)) returns

• Tensor –

  a \((0, 0)\) tensor: the predictor is never called, so the channel count

• Tensor –

  cannot be inferred.

Default test-time augmentation transforms for indoor semantic segmentation.

One view per (scale, rotation) pair, scales outermost, plus one x/y flip view when flip is set: the defaults give the 13-view precise-evaluation protocol of the ScanNet / ScanNet200 benchmarks. Rotations and the flip act on every key in keys (positions and normals); scales act on scale_keys only. Missing keys are skipped, so the same views serve scenes with and without normals. A rotation of \(0\) or a scale of \(1\) adds no transform.

Parameters:

• rotations (Sequence[float], default: (0.0, 90.0, 180.0, 270.0) ) –

  Rotation angles in degrees about axis, applied at every scale.

• scales (Sequence[float], default: (1.0, 0.95, 1.05) ) –

  Isotropic scale factors.

• flip (bool, default: True ) –

  Append one view mirrored along the x and y axes.

• axis (int, default: 2 ) –

  Rotation axis (2 is the up axis).

• keys (KeyCollection, default: (POS, NORMAL) ) –

  Keys rotated and flipped.

• scale_keys (KeyCollection, default: POS ) –

  Keys scaled.

Returns:

• List[Compose] –

  A list of len(rotations) * len(scales) + flip composed transforms for TTAInferer.

from torch_pointcloud.inferers import TTAInferer, VoxelPartitionInferer, simple_tta_transforms

inferer = TTAInferer(
    base=VoxelPartitionInferer(voxel_size=0.02, softmax=True, aggregate="sum"),
    transforms=simple_tta_transforms(),
)