From c11a1bb8263b278b2922564ba52007662603048e Mon Sep 17 00:00:00 2001 From: pushkalkatara Date: Fri, 8 May 2020 20:19:42 +0530 Subject: [PATCH] adding translate transform for point clouds Signed-off-by: pushkalkatara --- kaolin/transforms/pointcloudfunc.py | 97 +++++++++++++++++++++++++++++ kaolin/transforms/transforms.py | 58 +++++++++++++++++ 2 files changed, 155 insertions(+) diff --git a/kaolin/transforms/pointcloudfunc.py b/kaolin/transforms/pointcloudfunc.py index 4115b0995..ac90fcc74 100644 --- a/kaolin/transforms/pointcloudfunc.py +++ b/kaolin/transforms/pointcloudfunc.py @@ -25,6 +25,53 @@ EPS = 1e-6 +def shift(cloud: Union[torch.Tensor, PointCloud], + shf: Union[float, int, torch.Tensor], + inplace: Optional[bool] = True): + """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 (scaler, or tensor). + inplace (bool, optional): Bool to make the transform in-place + + Returns: + (torch.Tensor): shifted pointcloud pf 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.rand(1000,3) + >>> points2 = shift(points, torch.FloatTensor([3])) + """ + + if isinstance(cloud, np.ndarray): + cloud = torch.from_numpy(cloud) + + if isinstance(shf, np.ndarray): + shf = torch.from_numpy(shf) + + if isinstance(cloud, PointCloud): + cloud = cloud.points + + if isinstance(shf, int) or isinstance(shf, float): + shf = torch.Tensor([shf]).to(cloud.device) + + helpers._assert_tensor(cloud) + helpers._assert_tensor(shf) + helpers._assert_dim_ge(cloud, 2) + helpers._assert_gt(shf, 0.) + + if not inplace: + cloud = cloud.clone() + + return shf + cloud + + def scale(cloud: Union[torch.Tensor, PointCloud], scf: Union[float, int, torch.Tensor], inplace: Optional[bool] = True): @@ -74,6 +121,56 @@ def scale(cloud: Union[torch.Tensor, PointCloud], return scf * cloud +def translate(cloud: Union[torch.Tensor, PointCloud], tranmat: torch.Tensor, + inplace: Optional[bool] = True): + """Translate the input pointcloud by a translation matrix. + + Args: + cloud (Tensor or np.array): pointcloud (ndims = 2 or 3) + tranmat (Tensor or np.array): translation matrix (1 x 3, 1 per cloud). + + Returns: + cloud_tran (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.rand(1000,3) + >>> t_mat = torch.rand(1,3) + >>> points2 = translate(points, t_mat) + + """ + if isinstance(cloud, np.ndarray): + cloud = torch.from_numpy(cloud) + 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) + # Rotation matrix must have last two dimensions of shape 3. + helpers._assert_shape_eq(tranmat, (1, 3), dim=-1) + helpers._assert_shape_eq(tranmat, (1, 3), dim=-2) + + if not inplace: + cloud = cloud.clone() + + if tranmat.dim() == 2 and cloud.dim() == 2: + cloud = torch.add(tranmat, cloud) + else: + if tranmat.dim() == 2: + tranmat = tranmat.expand(cloud.shape[0], 1, 3) + cloud = torch.add(tranmat, cloud) + + return cloud + def rotate(cloud: Union[torch.Tensor, PointCloud], rotmat: torch.Tensor, inplace: Optional[bool] = True): """Rotates the the input pointcloud by a rotation matrix. diff --git a/kaolin/transforms/transforms.py b/kaolin/transforms/transforms.py index d687fb249..41ddeb55e 100644 --- a/kaolin/transforms/transforms.py +++ b/kaolin/transforms/transforms.py @@ -198,6 +198,35 @@ 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 pofactorint by which input + clouds are to be shifted. + inplace (bool, optional): Whether or not the transformation should be + in-place (default: True). + """ + + def __init__(self, shf: Union[int, float, torch.Tensor], + inplace: Optional[bool] = True): + self.shf = shf + self.inplace = inplace + + 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, inplace=self.inplace) + + 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 +260,35 @@ 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`). + inplace (bool, optional): Bool to make this operation in-place. + + TODO: Example. + + """ + + def __init__(self, tranmat: torch.Tensor, inplace: Optional[bool] = True): + self.tranmat = tranmat + self.inplace = inplace + + def __call__(self, cloud: Union[torch.Tensor, PointCloud]): + """ + Args: + cloud (torch.Tensor or PointCloud): Input pointcloud to be translated. + + Returns: + (torch.Tensor or PointCloud): Translated pointcloud. + """ + return pcfunc.translate(cloud, tranmat=self.tranmat, inplace=self.inplace) + + 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