geometry

Transforms that move points or derive geometric quantities.

Classes:

EstimateNormals –

Estimate per-point surface normals from coordinates via local PCA.
Shift –

Shift dictionary tensor entries by subtracting a computed offset.
BBoxCenter –

Derive the center of an axis-aligned bbox stored as a flat tensor.
Quantize –

Integer voxel-grid coordinates of every point, keeping the cloud at full resolution.
AxisMinOffset –

Per-point offset from a floor reference along a chosen coordinate axis.
Translate –

Add a constant offset to dictionary tensor entries.

EstimateNormals [source] ¶

EstimateNormals(
    keys: KeyCollection,
    dst_keys: KeyCollection = "normal",
    k: int = 16,
    orient_to_centroid: bool = False,
    batch_key: Optional[str] = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Estimate per-point surface normals from coordinates via local PCA.

Computes unit normals (see torch_pointcloud.transforms.functional.estimate_normals) for clouds that ship without them (e.g. S3DIS). Each normal is the least-variance direction of a point's \(k\) nearest neighbors. With orient_to_centroid, normals are flipped to face the cloud centroid.

Shift [source] ¶

Shift(
    keys: KeyCollection,
    method: ValueCollection[ShiftMethod],
    dim: int = 0,
    axes: Optional[Sequence[int]] = None,
    dst_keys: Optional[KeyCollection] = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Shift dictionary tensor entries by subtracting a computed offset.

Each key is offset independently. The offset is determined by method:

Method	Offset
`"bbox"`	Midrange: `(min + max) / 2`
`"centroid"`	Mean across the reduced dimension
`"min"`	Per-axis minimum (shifts to the positive octant)

On empty inputs (size \(0\) along dim) the tensor is returned unchanged.

method=centroidmethod=bboxmethod=minaxes

Shift method=centroid on an object

Shift method=bbox on an object

Shift method=min on an object

Shift restricted to a subset of axes, on an object

Parameters:

keys (KeyCollection) –

The keys to shift.
method (ValueCollection[ShiftMethod]) –

"bbox" (midrange), "centroid" (mean), or "min" (shift to origin).
dim (int, default: 0 ) –

The dimension to reduce over.
axes (Optional[Sequence[int]], default: None ) –

Which axes (last-dim indices) to shift. None (default) shifts every axis; pass e.g. axes=[0, 1] to recenter only XY. Axes outside this list are left unchanged - this is the composable knob for mixed-method shifts.
dst_keys (Optional[KeyCollection], default: None ) –

The keys to store the shifted data in.
allow_missing_keys (bool, default: False ) –

If True, skip missing keys silently.

Example

XY shifted by the bbox midpoint, Z shifted by its minimum (equivalent to the old CenterShift(apply_z=True)):

from torch_pointcloud.transforms import Compose, Shift

center_shift = Compose([
    Shift(keys="pos", method="bbox", axes=[0, 1]),  # XY: bbox midrange
    Shift(keys="pos", method="min",  axes=[2]),     # Z:  min
])

Without the Z step (equivalent to the old CenterShift(apply_z=False)), a single Shift suffices:

from torch_pointcloud.transforms import Shift

Shift(keys="pos", method="bbox", axes=[0, 1])

BBoxCenter [source] ¶

BBoxCenter(
    keys: KeyCollection,
    dst_keys: Optional[KeyCollection] = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Derive the center of an axis-aligned bbox stored as a flat tensor.

Reads a bbox at each source key, laid out as a \((2D,)\) vector \([\,\min_0, \ldots, \min_{D-1},\, \max_0, \ldots, \max_{D-1}\,]\), and writes the per-axis midpoint \((\min + \max) / 2\) (shape \((D,)\)) at the matching destination key.

ObjectScene

BBoxCenter on an object

BBoxCenter on a room

Parameters:

keys (KeyCollection) –

Source keys holding flat bbox tensors of shape \((2D,)\).
dst_keys (Optional[KeyCollection], default: None ) –

Destination keys for the centers. Defaults to overwriting the source keys.
allow_missing_keys (bool, default: False ) –

If True, silently skip absent source keys.

Example

import torch
from torch_pointcloud.transforms import BBoxCenter

data = {"block_bbox": torch.tensor([0.0, 0.0, 0.0, 1.5, 1.5, 2.8])}
BBoxCenter(keys="block_bbox", dst_keys="block_center")(data)
# data["block_center"] == tensor([0.75, 0.75, 1.40])

Quantize [source] ¶

Quantize(
    keys: KeyCollection,
    size: float,
    dst_keys: Optional[KeyCollection] = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Integer voxel-grid coordinates of every point, keeping the cloud at full resolution.

Stores \(\lfloor p / s \rfloor\) (shifted so the per-axis minimum is \(0\)) for each point of keys under dst_keys. Unlike Voxelize, no reduction happens: points sharing a voxel keep their own rows and get equal coordinates. This is how a voxel-partition evaluation feeds sparse models with every raw point (each sub-cloud holds one point per voxel, so its rows are exactly the voxels), and how test-time views recompute grid coordinates after rotating or scaling the positions.

ObjectScene

Quantize on an object

Quantize on a room

Parameters:

keys (KeyCollection) –

Keys holding point positions of shape \((N, D)\).
size (float) –

Voxel side length in the units of the positions.
dst_keys (Optional[KeyCollection], default: None ) –

Keys under which the grid coordinates are stored. Defaults to keys (in-place overwrite).
allow_missing_keys (bool, default: False ) –

If True, missing keys are skipped.

Example

import torch
from torch_pointcloud.transforms import Quantize

transform = Quantize(keys="pos", size=0.02, dst_keys="pos_grid")
data = transform({"pos": torch.tensor([[0.0, 0.0, 0.0], [0.03, 0.0, 0.0], [0.05, 0.0, 0.0]])})
data["pos_grid"]  # tensor([[0, 0, 0], [1, 0, 0], [2, 0, 0]])

AxisMinOffset [source] ¶

AxisMinOffset(
    keys: KeyCollection,
    axis: ValueCollection[int],
    quantile: Optional[float] = None,
    dst_keys: KeyCollection | None = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Per-point offset from a floor reference along a chosen coordinate axis.

For each point and a given axis \(a\) along tensor dimension \(d\), computes:

\[ o_i = p_{i,a} - r \]

where the floor reference \(r\) is either the strict minimum \(\min_j p_{j,a}\) (default) or, when quantile is set, the empirical quantile \(Q_q(p_{\cdot,a})\). A small positive quantile gives an outlier-robust floor estimate.

The result has the same shape as the input with the coordinate dimension reduced to size 1 (e.g. \((N, 3) \to (N, 1)\) or \((B, N, 3) \to (B, N, 1)\)). For batched inputs, the minimum is computed per-sample.

ObjectScene

AxisMinOffset on an object

AxisMinOffset on a room

Parameters:

keys (KeyCollection) –

Keys holding point positions of shape \((N, D)\).
axis (ValueCollection[int]) –

Coordinate axis \(a\) along which to compute the offset.
quantile (Optional[float], default: None ) –

Optional quantile \(q \in [0, 1]\) for the floor reference. When None, the strict per-axis minimum is used.
dst_keys (KeyCollection | None, default: None ) –

Keys under which the offset tensors are stored. Defaults to keys (in-place overwrite).
allow_missing_keys (bool, default: False ) –

If True, skip missing keys instead of raising.

Example

Let's say you have a point cloud with positions \((N, 3)\) in XYZ order and you want to compute the offset from the minimum along the z-axis, i.e. computing the height above the local floor.

import torch
from torch_pointcloud.transforms import AxisMinOffset

data = {
    "pos": torch.randn(10, 3),
}
transform = AxisMinOffset(keys="pos", dst_keys="pos_offset", axis=2)
data = transform(data)

Now, the data dictionary will contain the key pos_offset with the shape \((N, 1)\).

Translate [source] ¶

Translate(
    keys: KeyCollection,
    offset: Union[float, Sequence[float]],
    dst_keys: Optional[KeyCollection] = None,
    allow_missing_keys: bool = False,
)

Bases: DictTransform

Add a constant offset to dictionary tensor entries.

The deterministic counterpart of RandomTranslate and of Shift, whose offsets come from the data; after a Shift to the minimum, Translate(keys="pos", offset=(-0.75, -0.75, 0.0)) centers a 1.5 m block on its middle.

Parameters:

keys (KeyCollection) –

The keys to translate.
offset (Union[float, Sequence[float]]) –

Offset added to every row, a scalar or one value per channel of the last dimension.
dst_keys (Optional[KeyCollection], default: None ) –

The keys to store the translated data in.
allow_missing_keys (bool, default: False ) –

If True, skip missing keys silently.

Example

import torch
from torch_pointcloud.transforms import Translate
Translate(keys="pos", offset=(1.0, 0.0, -1.0))({"pos": torch.zeros(2, 3)})["pos"]  # tensor([[ 1.,  0., -1.],  # [ 1.,  0., -1.]])

estimate_normals [source] ¶

estimate_normals(
    pos: Tensor,
    k: int = 16,
    batch: Optional[Tensor] = None,
    orient_to_centroid: bool = False,
) -> Tensor

Estimate per-point unit surface normals by local PCA.

For each point the normal is the eigenvector of the smallest eigenvalue of the covariance of its \(k\) nearest neighbors, i.e. the direction of least variance (the local tangent-plane normal).

PCA gives no canonical orientation. By default the sign is the arbitrary-but-deterministic sign returned by torch.linalg.eigh. With orient_to_centroid, each normal is flipped to point towards its cloud's centroid, which approximates the inward-facing orientation of meshes scanned from inside a room (S3DIS, ScanNet) and matters when the consuming model was trained on oriented normals.

Parameters:

pos (Tensor) –

Point coordinates of shape \((N, 3)\).
k (int, default: 16 ) –

Number of nearest neighbors (the point itself included) used per local PCA. Must not exceed the number of points in the smallest cloud.
batch (Optional[Tensor], default: None ) –

Optional \((N,)\) batch index so neighbors never cross cloud boundaries.
orient_to_centroid (bool, default: False ) –

If True, flip each normal to point towards its cloud's centroid.

Returns:

Tensor –

Unit normals of shape \((N, 3)\).

Raises:

ValueError –

If pos has fewer than k points.

Shape

Input: \((N, 3)\)
Output: \((N, 3)\)

shift [source] ¶

shift(
    x: Tensor,
    method: ShiftMethod,
    dim: int = 0,
    axes: Optional[Sequence[int]] = None,
) -> Tensor

Subtract a data-driven offset from x.

The offset is computed from x itself along the reduction dimension dim:

`method`	Offset
`"bbox"`	Midrange: `(min + max) / 2`
`"centroid"`	Mean across the reduced dimension
`"min"`	Per-axis minimum (shifts to the positive octant)

When axes is given, only those axis-indices of the offset are non-zero, so axes not listed are left untouched. This is the composable knob for mixed-method shifts:

import torch

import torch_pointcloud.transforms.functional as F

x = torch.rand(100, 3)
# Center XY at the bbox midpoint and Z at the minimum
x = F.shift(x, method="bbox", axes=[0, 1])
x = F.shift(x, method="min", axes=[2])

The two calls touch disjoint axes, so they commute.

Parameters:

x (Tensor) –

Input tensor.
method (ShiftMethod) –

How the offset is computed. See the table.
dim (int, default: 0 ) –

The dimension to reduce over when computing the offset.
axes (Optional[Sequence[int]], default: None ) –

Last-dim axis indices to shift. None (default) shifts every axis.

Returns:

Tensor –

The shifted tensor, same shape as x. Returns x unchanged when
Tensor –

x.size(dim) == 0.

Raises:

ValueError –

If method is not one of "bbox", "centroid", "min".

quantize [source] ¶

quantize(pos: Tensor, size: float) -> Tensor

Integer voxel-grid coordinates of every point, without reducing the cloud.

Each point maps to \(\lfloor p / s \rfloor\) shifted so the per-axis minimum is \(0\); points sharing a voxel get equal coordinates and every input row is kept. This is the coordinate a voxel-partition protocol feeds to a sparse model for each raw point (Voxelize(pos_reduce="grid") produces the same coordinates for the one representative it keeps per voxel).

Parameters:

pos (Tensor) –

Point positions of shape \((N, D)\).
size (float) –

Voxel side length in the units of pos.

Returns:

Tensor –

Long tensor of shape \((N, D)\) (empty input returns an empty \((0, D)\) tensor).

Example

import torch
from torch_pointcloud.transforms import functional as F

pos = torch.tensor([[0.0, 0.0, 0.0], [0.03, 0.0, 0.0], [0.05, 0.0, 0.0]])
F.quantize(pos, size=0.02)  # tensor([[0, 0, 0], [1, 0, 0], [2, 0, 0]])

axis_min_offset [source] ¶

axis_min_offset(
    x: Tensor, axis: int, quantile: Optional[float] = None
) -> Tensor

Per-point offset from a floor reference along a chosen coordinate axis.

For positions of shape \((N, D)\) and an axis \(a \in [0, D)\), returns a tensor of shape \((N, 1)\) whose entries are \(x_{i, a} - r\) where the floor reference \(r\) is either the strict minimum \(\min_j x_{j, a}\) (default) or, when quantile is given, the empirical quantile \(Q_{q}(x_{\cdot, a})\). A small positive quantile (e.g. \(q = 0.0099\)) yields an outlier-robust floor estimate. Useful for extracting "height above the local floor" as a per-point feature.

Parameters:

x (Tensor) –

Input tensor of shape \((N, D)\).
axis (int) –

Axis index in the last dimension.
quantile (Optional[float], default: None ) –

Optional quantile \(q \in [0, 1]\) for the floor reference. When None, the strict per-axis minimum is used (equivalent to \(q = 0\)).

Returns:

Tensor –

Tensor of shape \((N, 1)\) with the same dtype as x. Returns an empty
Tensor –

\((0, 1)\) tensor when x is empty.

rotation_matrix [source] ¶

rotation_matrix(
    angle: float,
    axis: int = 2,
    device: Optional[device] = None,
) -> Tensor

\(3 \times 3\) rotation matrix for angle radians around an axis-aligned axis.

Parameters:

angle (float) –

Rotation angle in radians.
axis (int, default: 2 ) –

Axis index to rotate around (0=X, 1=Y, 2=Z).
device (Optional[device], default: None ) –

Output device. Defaults to CPU.

Returns:

Tensor –

Rotation matrix of shape \((3, 3)\).

Raises:

ValueError –

If axis is not in {0, 1, 2}.

rotate_vectors [source] ¶

rotate_vectors(x: Tensor, rotation: Tensor) -> Tensor

Rotate a packed field of 3D vectors by a rotation matrix.

Each contiguous triple of the last dimension rotates as a vector, so it handles both a plain \((N, 3)\) field (e.g. coordinates or normals) and a \((N, 3 G)\) field of \(G\) tiled offsets (e.g. vote offsets) alike.

Parameters:

x (Tensor) –

Vector field of shape \((N, 3)\) or \((N, 3 G)\).
rotation (Tensor) –

A \(3 \times 3\) rotation matrix.

Returns:

Tensor –

The rotated tensor with the same shape as x.