tta

Test-time augmentation (TTA) inferer.

Wraps any Inferer and runs it N times, each time under a different spatial augmentation of the input, then aggregates the per-point predictions. Because point cloud segmentation outputs are indexed by point ID rather than by spatial position, predictions from rotated or flipped views are already aligned and can be averaged directly without inverting the transform.

Classes:

TTAInferer –

Test-time augmentation inferer.

Functions:

simple_tta_transforms –

Default test-time augmentation transforms for indoor semantic segmentation.

TTAInferer ¶

TTAInferer(
    base: Inferer,
    transforms: Union[TransformFn, Sequence[TransformFn]],
    num_passes: Optional[int] = None,
    include_identity: bool = False,
    aggregate: AggregateMode = "mean",
    ema_smoothing: float = 0.95,
    ema_softmax: bool = True,
    pos_key: str = POS,
)

Bases: Inferer

Test-time augmentation inferer.

Runs the wrapped base inferer once per augmentation pass and aggregates the per-point predictions across passes.

Two augmentation modes are supported:

Single callable: re-sampled independently each pass. Use for random augmentations such as uniformly random rotation. Requires num_passes.
Sequence of callables: each element is applied to exactly one pass in order. Use for a fixed view set (e.g. 8 evenly-spaced rotations). num_passes is inferred from the sequence length.

Parameters:

base (Inferer) –

Underlying Inferer invoked once per pass. Any concrete inferer works: SimpleInferer(), SlidingWindowInferer(...), KNNWindowInferer(...).
transforms (Union[TransformFn, Sequence[TransformFn]]) –

Single callable (re-sampled each pass) or a sequence of callables for fixed views. Any Dict[str, Any] -> Dict[str, Any] callable works, including Compose. When a sequence is given, num_passes is ignored.
num_passes (Optional[int], default: None ) –

Number of TTA passes when transforms is a single callable. Must be \(\geq 1\).
include_identity (bool, default: False ) –

If True, run one extra pass on the un-augmented input before the augmented passes (the "clean + N random views" voting protocol), so the total pass count is num_passes + 1.
aggregate (AggregateMode, default: 'mean' ) –

How per-point predictions are combined across passes. "mean" averages per-pass outputs (works for logits or probabilities). "ema" maintains an exponential moving average of softmax probabilities.
ema_smoothing (float, default: 0.95 ) –

EMA factor \(\alpha \in [0, 1)\) used when aggregate="ema".
ema_softmax (bool, default: True ) –

When aggregate="ema", softmax each pass's output before accumulating. Set False if the base inferer already returns probabilities (e.g. SlidingWindowInferer(softmax=True)).
pos_key (str, default: POS ) –

Dict key for the position tensor (used for the empty-output fallback).

Example

A 4-pass TTA over random Z rotations and X/Y flips:

from torch_pointcloud.inferers import TTAInferer, SlidingWindowInferer
from torch_pointcloud.transforms import Compose, RandomRotate, RandomFlip

base = SlidingWindowInferer(block_size=6.0)
aug = Compose([
    RandomRotate(keys="pos", angle_range=(-180.0, 180.0), axis=2, p=1.0),
    RandomFlip(keys="pos", axes=[0, 1], p=0.5),
])
inferer = TTAInferer(base=base, transforms=aug, num_passes=4,
                     aggregate="mean")
probs = inferer(data, predictor=lambda d: model(d["pos"], d["pos"], d["batch"]))

Enumerated 8-view TTA (ScanNet ablation):

views = [Compose([RandomRotate(keys="pos", angle_range=(a, a), axis=2, p=1.0)])
         for a in (0.0, 45.0, 90.0, 135.0, 180.0, 225.0, 270.0, 315.0)]
inferer = TTAInferer(base=base, transforms=views, aggregate="mean")

simple_tta_transforms ¶

simple_tta_transforms(
    rotations: Sequence[float] = (0.0, 90.0, 180.0, 270.0),
    scales: Sequence[float] = (1.0, 0.95, 1.05),
    flip: bool = True,
    axis: int = 2,
    keys: KeyCollection = (POS, NORMAL),
    scale_keys: KeyCollection = POS,
) -> List[Compose]

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.

Example

from torch_pointcloud.inferers import TTAInferer, VoxelPartitionInferer, simple_tta_transforms

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