knn_window

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

Implements a coverage-driven iterative loop: each step selects the least-covered point as a window center, crops its \(k\) nearest neighbors, runs the predictor on that crop, and accumulates per-point predictions weighted by distance to the center. The loop ends once every point's coverage score exceeds a threshold.

Windows adapt to point density and naturally prioritize under-covered regions. Because windows overlap, each point is typically predicted several times; overlapping predictions are combined by a weighted mean or an exponential moving average (EMA).

Classes:

KNNWindowInferer –

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

Functions:

knn_window_inference –

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

KNNWindowInferer ¶

KNNWindowInferer(
    roi_num_points: int = 65536,
    sw_batch_size: int = 1,
    overlap: float = 0.5,
    mode: WindowMode = "constant",
    sigma_scale: float = 0.125,
    aggregate: AggregateMode = "weighted_mean",
    ema_smoothing: float = 0.95,
    softmax: bool = False,
    transform: Optional[
        Callable[[Dict[str, Any]], Dict[str, Any]]
    ] = None,
    pos_key: str = POS,
    batch_key: str = BATCH,
    progress: bool = False,
    seed: Optional[int] = None,
)

Bases: Inferer

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

Maintains a per-point coverage score and iteratively crops KNN windows around 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.

Example

from torch_pointcloud.inferers import KNNWindowInferer

# EMA aggregation: outputs calibrated probabilities directly.
inferer = KNNWindowInferer(roi_num_points=65_536, overlap=0.5, aggregate="ema")
probs = inferer(data, predictor=lambda d: model(d["pos"], d["pos"], d["batch"]))

knn_window_inference ¶

knn_window_inference(
    data: Dict[str, Any],
    *,
    predictor: Callable[[Dict[str, Any]], Tensor],
    roi_num_points: int = 65536,
    sw_batch_size: int = 1,
    overlap: float = 0.5,
    mode: WindowMode = "constant",
    sigma_scale: float = 0.125,
    aggregate: AggregateMode = "weighted_mean",
    ema_smoothing: float = 0.95,
    softmax: bool = False,
    transform: Optional[
        Callable[[Dict[str, Any]], Dict[str, Any]]
    ] = None,
    pos_key: str = POS,
    batch_key: str = BATCH,
    progress: bool = False,
    seed: Optional[int] = None,
) -> Tensor

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:

"weighted_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="weighted_mean"; EMA updates ignore distance weights.
sigma_scale (float, default: 0.125 ) –

Gaussian sigma scale factor (only used when mode="gaussian").
aggregate (AggregateMode, default: 'weighted_mean' ) –

How predictions from overlapping windows are combined. "weighted_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")).
pos_key (str, default: POS ) –

Dict key for the position tensor.
batch_key (str, default: BATCH ) –

Dict key for the per-point batch index.
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.

Returns:

Tensor –

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

logits when aggregate="weighted_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.