aidsorb.transforms.points#

Helper functions and classes for transforming point clouds.

Note

All transforms expect an input Tensor of shape (N, 3+C).

class aidsorb.transforms.points.Center[source]#

Bases: object

Center the coordinates of a point cloud by subtracting their centroid.

See also

center_pcd()

For a functional interface.

Examples

>>> x = torch.arange(4.)
>>> pcd = torch.stack((x, x))
>>> center = Center()
>>> center(pcd)
tensor([[0., 0., 0., 3.],
        [0., 0., 0., 3.]])
class aidsorb.transforms.points.RandomErase(n_points, local=False)[source]#

Bases: object

Randomly erase points from the point cloud.

Parameters:

  • n_points (int or float) – Number or fraction of points to be erased. If float, it shoud be in the interval (0, 1). In this case, int(len(pcd) * n_points) points are erased.

  • local (bool, default=False) – Whether to erase a local or global patch of n_points.

Examples

>>> pcd = torch.randn(100, 5)
>>> erase = RandomErase(n_points=10)
>>> erase(pcd).shape
torch.Size([90, 5])

>>> # Erase a global patch.
>>> pcd = torch.randn(100, 5)
>>> erase = RandomErase(n_points=0.4)
>>> erase(pcd).shape
torch.Size([60, 5])

>>> # Erase a local patch.
>>> pcd = torch.randn(50, 4)
>>> erase = RandomErase(n_points=0.7, local=True)
>>> erase(pcd).shape
torch.Size([15, 4])

>>> pcd = torch.randn(100, 5)
>>> erase = RandomErase(n_points=100)
>>> erase(pcd)
Traceback (most recent call last):
    ...
RuntimeError: resulting point cloud has no points

>>> pcd = torch.randn(100, 5)
>>> erase = RandomErase(n_points=150, local=True)
>>> erase(pcd)
Traceback (most recent call last):
    ...
RuntimeError: resulting point cloud has no points
class aidsorb.transforms.points.RandomFlip[source]#

Bases: object

Flip the coordinates of a point cloud along a randomly selected axis.

Notes

The input tensor is copied to prevent in-place modifications and preserve the original data.

Examples

>>> pcd = torch.randn(10, 5)
>>> flip = RandomFlip()
>>> new_pcd = flip(pcd)
>>> new_pcd.shape
torch.Size([10, 5])

>>> coords, feats = split_pcd(pcd)
>>> new_coords, new_feats = split_pcd(new_pcd)
>>> torch.equal(coords, new_coords)  # Coordinates are affected.
False
>>> torch.equal(feats, new_feats)  # Features are not affected.
True

>>> # Only one axis is flipped.
>>> (pcd == -new_pcd).all(0).sum()
tensor(1)
class aidsorb.transforms.points.RandomJitter(std, n_points=None, local=None)[source]#

Bases: object

Jitter the coordinates of a point cloud by adding zero-mean normal noise.

  • If both n_points and local are None, then all points are jittered.

  • If local!=None, then n_points must be specified.

Parameters:

  • std (float) – Standard deviation of the normal noise.

  • n_points (int or float or None, default=None) – Number or fraction of points to be jittered.

  • local (bool or None, default=None) – Whether to jitter a local or global patch of n_points.

Examples

>>> # Jitter all points.
>>> pcd = torch.randn(100, 5)
>>> jitter = RandomJitter(0.1)
>>> new_pcd = jitter(pcd)
>>> new_pcd.shape
torch.Size([100, 5])

>>> coords, feats = split_pcd(pcd)
>>> new_coords, new_feats = split_pcd(new_pcd)
>>> torch.equal(new_coords, coords)  # Coordinates are affected.
False
>>> torch.equal(new_feats, feats)  # Features are not affected.
True

>>> # Jitter a subset of points.
>>> pcd = torch.randn(30, 4)
>>> jitter =  RandomJitter(0.5, n_points=0.3, local=True)
>>> new_pcd = jitter(pcd)
>>> new_pcd.shape
torch.Size([30, 4])

>>> coords, feats = split_pcd(pcd)
>>> new_coords, new_feats = split_pcd(new_pcd)
>>> torch.equal(new_coords, coords)  # Coordinates are affected.
False
>>> torch.equal(new_feats, feats)  # Features are not affected.
True

>>> (new_pcd == pcd).all(1).sum()
tensor(21)
class aidsorb.transforms.points.RandomRotation[source]#

Bases: object

Randomly rotate the coordinates of a point cloud.

Examples

>>> pcd = torch.randn(25, 4)
>>> rot = RandomRotation()
>>> new_pcd = rot(pcd)
>>> new_pcd.shape
torch.Size([25, 4])

>>> coords, feats = split_pcd(pcd)
>>> new_coords, new_feats = split_pcd(new_pcd)

>>> torch.equal(new_coords, coords)  # Coordinates are affected.
False
>>> torch.equal(new_feats, feats)  # Features are not affected.
True
class aidsorb.transforms.points.RandomSample(size)[source]#

Bases: object

Sample without replacement a number of points from the point cloud.

If size >= len(pcd), the original point cloud is returned unchanged.

Parameters:

size (int) – Number of points to sample.

Examples

>>> pcd = torch.randn(10, 4)
>>> sample = RandomSample(size=5)
>>> sample(pcd).shape
torch.Size([5, 4])

>>> # No sampling, original point cloud is returned.
>>> pcd = torch.randn(10, 4)
>>> sample = RandomSample(size=100)
>>> torch.equal(pcd, sample(pcd))
True
aidsorb.transforms.points.center_pcd(pcd)[source]#

Center the coordinates of a point cloud by subtracting their centroid.

Parameters:

pcd (tensor of shape (N, 3+C))

Returns:

new_pcd – Centered point cloud.

Return type:

tensor of shape (N, 3+C)

Examples

>>> pcd = torch.tensor([[2., 1., 3., 6.], [-3., 2., 8., 8.]])
>>> new_pcd = center_pcd(pcd)
>>> new_pcd.mean(dim=0)
tensor([0., 0., 0., 7.])

>>> pcd = torch.randn(5, 2)  # Invalid shape.
>>> new_pcd = center_pcd(pcd)
Traceback (most recent call last):
    ...
ValueError: expecting shape (N, 3+C) but received shape (5, 2)
aidsorb.transforms.points.split_pcd(pcd)[source]#

Split a point cloud to coordinates and features.

Parameters:

pcd (tensor of shape (N, 3+C))

Returns:

coords_feats – Coordinates and features of point cloud as (coords, feats).

  • coords tensor of shape (N, 3)

  • feats tensor of shape (N, C)

Return type:

tuple

Examples

>>> pcd = torch.randn(25, 7)  # Point cloud with 4 features.
>>> coords, feats = split_pcd(pcd)
>>> coords.shape
torch.Size([25, 3])
>>> feats.shape
torch.Size([25, 4])

>>> pcd = torch.randn(15, 3)  # Point cloud with no features.
>>> coords, feats = split_pcd(pcd)
>>> coords.shape
torch.Size([15, 3])
>>> feats.shape
torch.Size([15, 0])
aidsorb.transforms.points.transform_pcd(pcd, tfm)[source]#

Transform the coordinates of a point cloud.

For molecular point clouds, only rigid transformations are recommended.

Parameters:

  • pcd (tensor of shape (N, 3+C)) – Original point cloud.

  • tfm (tensor of shape (3, 3)) – Transformation matrix.

Returns:

new_pcd – Transformed point cloud.

Return type:

tensor of shape (N, 3+C)

Examples

>>> pcd = torch.tensor([[3, -9, 2, 6], [3, 4, -1, 8]])
>>> tfm = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]])
>>> transform_pcd(pcd, tfm)
tensor([[ 9,  3,  2,  6],
        [-4,  3, -1,  8]])

>>> pcd = torch.randn(424, 2)  # Invalid shape.
>>> transform_pcd(pcd, tfm)
Traceback (most recent call last):
    ...
ValueError: expecting shape (N, 3+C) but received shape (424, 2)
aidsorb.transforms.points.upsample_pcd(pcd, size)[source]#

Upsample pcd to a new size by sampling with replacement from pcd.

Parameters:

  • pcd (tensor of shape (N, C)) – Original point cloud of size N.

  • size (int) – Size of the new point cloud.

Returns:

new_pcd

Return type:

tensor of shape (size, C)

Examples

>>> pcd = torch.tensor([[2, 4, 5, 6]])
>>> upsample_pcd(pcd, 3)
tensor([[2, 4, 5, 6],
        [2, 4, 5, 6],
        [2, 4, 5, 6]])

>>> # New points must be from pcd.
>>> pcd = torch.randn(10, 4)
>>> new_pcd = upsample_pcd(pcd, 20)
>>> (new_pcd[-1] == pcd).all(1).any()  # Check for last point.
tensor(True)

>>> # New size must be greater than the original.
>>> pcd = torch.randn(10, 4)
>>> new_pcd = upsample_pcd(pcd, 5)
Traceback (most recent call last):
    ...
ValueError: target size (5) must be greater than the original size (10)