Translate transform for pointclouds

kaolin/transforms/pointcloudfunc.py

@@ -25,6 +25,42 @@
 EPS = 1e-6
 
 
+def shift(cloud: Union[torch.Tensor, PointCloud],
+          shf: Union[float, int, torch.Tensor]):
+    """Shift the input pointcloud by a shift factor.
+
+    Args:
+        cloud (torch.Tensor or kaolin.rep.PointCloud): pointcloud (ndims >= 2).
+        shf (float, int, torch.Tensor): shift factor (scalar, or tensor).
+
+    Returns:
+        (torch.Tensor): shifted pointcloud of the same shape as input.
+
+    Shape:
+        - cloud: :math:`(B x N x D)` (or) :math:`(N x D)`, where :math:`(B)`
+            is the batchsize, :math:`(N)` is the number of points per cloud,
+            and :math:`(D)` is the dimensionality of each cloud.
+        - shf: :math:`(1)` or :math:`(B)`.
+
+    Example:
+        >>> points = torch.tensor([[0., 0., 0.], 
+                                  [1., 2., 3.],
+                                  [2., 3., 4.]])
+        >>> points2 = shift(points, 3)
+    """
+    if isinstance(cloud, PointCloud):
+        cloud = cloud.points
+
+    if isinstance(shf, int) or isinstance(shf, float):
+        shf = torch.Tensor([shf], device=cloud.device)
+
+    helpers._assert_tensor(cloud)
+    helpers._assert_tensor(shf)
+    helpers._assert_dim_ge(cloud, 2)
+
+    return shf + cloud
+
+
 def scale(cloud: Union[torch.Tensor, PointCloud],
           scf: Union[float, int, torch.Tensor],
           inplace: Optional[bool] = True):
@@ -50,18 +86,11 @@ def scale(cloud: Union[torch.Tensor, PointCloud],
         >>> points2 = scale(points, torch.FloatTensor([3]))
 
     """
-
-    if isinstance(cloud, np.ndarray):
-        cloud = torch.from_numpy(cloud)
-
-    if isinstance(scf, np.ndarray):
-        scf = torch.from_numpy(scf)
-
     if isinstance(cloud, PointCloud):
         cloud = cloud.points
 
     if isinstance(scf, int) or isinstance(scf, float):
-        scf = torch.Tensor([scf]).to(cloud.device)
+        scf = torch.Tensor([scf], device=cloud.device)
 
     helpers._assert_tensor(cloud)
     helpers._assert_tensor(scf)
@@ -74,17 +103,57 @@ def scale(cloud: Union[torch.Tensor, PointCloud],
     return scf * cloud
 
 
+def translate(cloud: Union[torch.Tensor, PointCloud], tranmat: torch.Tensor):
+    """Translate the input pointcloud by a translation matrix.
+
+    Args:
+        cloud (torch.Tensor or kaolin.rep.PointCloud): pointcloud (ndims = 2 or 3)
+        tranmat (torch.Tensor or np.array): translation matrix (1 x 3, 1 per cloud).
+
+    Returns:
+        (torch.Tensor): Translated pointcloud of the same shape as input.
+
+    Shape:
+        - cloud: :math:`(B x N x 3)` (or) :math:`(N x 3)`, where :math:`(B)`
+            is the batchsize, :math:`(N)` is the number of points per cloud,
+            and :math:`(3)` is the dimensionality of each cloud.
+        - tranmat: :math:`(1, 3)` or :math:`(B, 1, 3)`.
+
+    Example:
+        >>> points = torch.tensor([[0., 0., 0.], 
+                                  [1., 2., 3.],
+                                  [2., 3., 4.]])
+        >>> t_mat = torch.rand(1,3)
+        >>> points2 = translate(points, t_mat)
+
+    """
+    if isinstance(cloud, PointCloud):
+        cloud = cloud.points
+    if isinstance(tranmat, np.ndarray):
+        trainmat = torch.from_numpy(tranmat)
+
+    helpers._assert_tensor(cloud)
+    helpers._assert_tensor(tranmat)
+    helpers._assert_dim_ge(cloud, 2)
+    helpers._assert_dim_ge(tranmat, 2)
+    helpers._assert_dim_le(cloud, 3)
+    helpers._assert_dim_le(tranmat, 3)
+    # Translation matrix and cloud last dimension must be equal
+    helpers._assert_shape_eq(tranmat, cloud.shape, dim=-1)
+
+    return torch.add(tranmat, cloud)
+
 def rotate(cloud: Union[torch.Tensor, PointCloud], rotmat: torch.Tensor,
            inplace: Optional[bool] = True):
     """Rotates the the input pointcloud by a rotation matrix.
 
     Args:
-        cloud (Tensor or np.array): pointcloud (ndims = 2 or 3)
-        rotmat (Tensor or np.array): rotation matrix (3 x 3, 1 per cloud).
+        cloud (torch.Tensor or kaolin.rep.PointCloud): pointcloud (ndims = 2 or 3)
+        rotmat (torch.Tensor or np.array): rotation matrix (3 x 3, 1 per cloud).
         inplace (bool, optional): Bool to make the transform in-place.
 
     Returns:
-        cloud_rot (Tensor): rotated pointcloud of the same shape as input
+        (torch.Tensor): rotated pointcloud of the same shape as input
 
     Shape:
         - cloud: :math:`(B x N x 3)` (or) :math:`(N x 3)`, where :math:`(B)`
@@ -98,8 +167,6 @@ def rotate(cloud: Union[torch.Tensor, PointCloud], rotmat: torch.Tensor,
         >>> points2 = rotate(points, r_mat)
 
     """
-    if isinstance(cloud, np.ndarray):
-        cloud = torch.from_numpy(cloud)
     if isinstance(cloud, PointCloud):
         cloud = cloud.points
     if isinstance(rotmat, np.ndarray):
@@ -133,15 +200,15 @@ def realign(src: Union[torch.Tensor, PointCloud],
     box as that of pointcloud `tgt`.
 
     Args:
-        src (torch.Tensor or PointCloud) : Source pointcloud to be transformed
-            (shape: :math:`\cdots \times N \times D`, where :math:`N` is the
-            number of points in the pointcloud, and :math:`D` is the
-            dimensionality of each point in the cloud).
-        tgt (torch.Tensor or PointCloud) : Target pointcloud to which `src`is
-            to be transformed (The `src` cloud is transformed to the
-            axis-aligned bounding box that the target cloud maps to). This
-            cloud must have the same number of dimensions :math:`D` as in the
-            source cloud. (shape: :math:`\cdots \times \cdots \times D`).
+        src ((torch.Tensor or kaolin.rep.PointCloud)) : Source pointcloud to
+             be transformed (shape: :math:`\cdots \times N \times D`, where 
+             :math:`N` is the number of points in the pointcloud, and :math:`D`
+             is the dimensionality of each point in the cloud).
+        tgt ((torch.Tensor or kaolin.rep.PointCloud)) : Target pointcloud to 
+            which `src`is to be transformed (The `src` cloud is transformed 
+            to the axis-aligned bounding box that the target cloud maps to). 
+            This cloud must have the same number of dimensions :math:`D` as 
+            in the source cloud. (shape: :math:`\cdots \times \cdots \times D`).
         inplace (bool, optional): Bool to make the transform in-place.
 
     Returns:
@@ -183,20 +250,17 @@ def normalize(cloud: Union[torch.Tensor, PointCloud],
     deviation. For batched clouds, each cloud is independently normalized.
 
     Args:
-        cloud (torch.Tensor or PointCloud): Input pointcloud to be normalized
-            (shape: :math:`B \times \cdots \times N \times D`, where :math:`B`
-            is the batchsize (optional), :math:`N` is the number of points in
-            the cloud, and :math:`D` is the dimensionality of the cloud.
+        cloud (torch.Tensor or kaolin.rep.PointCloud): Input pointcloud to
+            be normalized (shape: :math:`B \times \cdots \times N \times D`, 
+            where :math:`B` is the batchsize (optional), :math:`N` is the 
+            number of points in the cloud, and :math:`D` is the dimensionality
+            of the cloud. 
         inplace (bool, optional): Bool to make the transform in-place.
 
     Returns:
-        (torch.Tensor or PointCloud): The normalized pointcloud.
+        (torch.Tensor or kaolin.rep.PointCloud): The normalized pointcloud.
 
     """
-
-    if isinstance(cloud, np.ndarray):
-        cloud = torch.from_numpy(cloud)
-
     helpers._assert_tensor(cloud)
     helpers._assert_dim_ge(cloud, 2)
     if not inplace:

kaolin/transforms/transforms.py

@@ -198,6 +198,31 @@ def __call__(self, arr: np.ndarray):
         return torch.from_numpy(arr)
 
 
+class ShiftPointCloud(Transform):
+    r"""Shift a pointcloud with respect a fixed shift factor.
+    Given a shift factor `shf`, this transform will shift each point in the 
+    pointcloud, i.e.,
+    ``cloud = shf + cloud``
+
+    Args:
+        shf (int or float or torch.Tensor): Shift factor by which input
+            clouds are to be shifted.
+    """
+
+    def __init__(self, shf: Union[int, float, torch.Tensor]):
+        self.shf = shf
+
+    def __call__(self, cloud: Union[torch.Tensor, PointCloud]):
+        """
+        Args:
+            cloud (torch.Tensor or PointCloud): Pointcloud to be shifted.
+
+        Returns:
+            (torch.Tensor or PointCloud): Shifted pointcloud.
+        """
+        return pcfunc.shift(cloud, shf=self.shf)
+
+
 class ScalePointCloud(Transform):
     """Scale a pointcloud with a fixed scaling factor.
     Given a scale factor `scf`, this transform will scale each point in the
@@ -231,6 +256,38 @@ def __call__(self, cloud: Union[torch.Tensor, PointCloud]):
         return pcfunc.scale(cloud, scf=self.scf, inplace=self.inplace)
 
 
+class TranslatePointCloud(Transform):
+    r"""Translate a pointcloud with a given translation matrix.
+    Given a :math:`1 \times 3` translation matrix, this transform will 
+    translate each point in the cloud by the translation matrix specified.
+
+    Args:
+        tranmat (torch.Tensor): Translation matrix that specifies the translation 
+            to be applied to the pointcloud (shape: :math:`1 \times 3`).
+
+    Example:
+        import kaolin.transforms as tfs
+        tranmat = torch.ones(1,3)
+        translate_fn = tfs.TranslatePointCloud(tranmat)
+        pc = torch.rand(1000,3)
+        translated_pc = translate_fn(pc)
+    """
+
+    def __init__(self, tranmat: torch.Tensor):
+        self.tranmat = tranmat
+
+    def __call__(self, cloud: Union[torch.Tensor, PointCloud]):
+        """
+        Args:
+            cloud (torch.Tensor or kaolin.rep.PointCloud): Input pointcloud 
+            to be translated.
+
+        Returns:
+            (torch.Tensor or kaolin.rep.PointCloud): Translated pointcloud.
+        """
+        return pcfunc.translate(cloud, tranmat=self.tranmat)
+
+
 class RotatePointCloud(Transform):
     r"""Rotate a pointcloud with a given rotation matrix.
     Given a :math:`3 \times 3` rotation matrix, this transform will rotate each

tests/transforms/test_transforms.py

@@ -30,6 +30,13 @@ def test_numpy_to_tensor(device='cpu'):
     assert torch.is_tensor(ten)
 
 
+def test_shift_pointcloud(device='cpu'):
+    unit_shift = kal.transforms.ShiftPointCloud(-1)
+    pc = torch.ones(4, 3)
+    pc_ = unit_shift(pc)
+    assert_allclose(pc_, torch.zeros(4, 3))
+
+
 def test_scale_pointcloud(device='cpu'):
     twice = kal.transforms.ScalePointCloud(2)
     halve = kal.transforms.ScalePointCloud(0.5)
@@ -38,6 +45,19 @@ def test_scale_pointcloud(device='cpu'):
     assert_allclose(pc_, torch.ones(4, 3))
 
 
+def test_translate_pointcloud(device='cpu'):
+    pc = torch.ones(4, 3)
+    tmat = torch.tensor([[-1.0,-1.0,-1.0]])
+    translate = kal.transforms.TranslatePointCloud(tmat)
+    pc_ = translate(pc)
+    assert_allclose(pc_, torch.zeros(4, 3))
+    pc = torch.ones(4, 2)
+    tmat = torch.tensor([[-1.0,-1.0]])
+    translate = kal.transforms.TranslatePointCloud(tmat)
+    pc_ = translate(pc)
+    assert_allclose(pc_, torch.zeros(4, 2))
+
+
 def test_rotate_pointcloud(device='cpu'):
     pc = torch.ones(4, 3)
     rmat = 2 * torch.eye(3)