How do I rotate sparse pointclouds using pytorch3d/the gpu?

[PottedRosePetal]

So I have pointclouds with the shape (batch_size, num_points, coordinates). I now want to rotate them around each axis by angles in another tensor with the shape (batch_size, angles). This is normally (192,2048,3) for the pointclouds and (192, 3) for the angles.

I already implemented this using the open3D library, however, it seems like the rotation function is not supported on the gpu. Since this is very slow and I want live data augmentation, I am cpu bottlenecked. During the search for some other stuff, I found that pytorch3d probably supports the same rotation, and I assume with tensors and on the gpu. I have looked through the docs and some questions, but I couldnt find some direct hint on how to do this? The rotate functions are not clearly defined and I am not sure if i should use quaternion operations, some pointcloud object... honestly I am kinda overwhelmed.

So far I have done it this way, just so the actual use is clear:

    def apply_augmentation(pointcloud:torch.Tensor, aug_args:dict) -> np.ndarray:
        pcd = o3d.t.geometry.PointCloud()
        pcd.point.positions = o3c.Tensor(pointcloud.numpy())
        center = o3c.Tensor(torch.mean(pointcloud, dim=0).numpy())
        rotation_matrix = o3d.geometry.get_rotation_matrix_from_xyz(aug_args["rotation"])
        pcd.rotate(rotation_matrix, center=center)

        return pcd.point.positions.numpy()

I hope you can help me, this would save me a ton of time! I need to run this over 100k iterations every few days on not so great hardware :)

[bottler]

What format are the angles? Maybe you need something like this, which can happen with all data on gpu

from pytorch3d.transforms import euler_angles_to_matrix
matrices = euler_angles_to_matrix(angles, some_convention)
pointclouds2 = pointclouds@matrices

[PottedRosePetal]

oh perfect, thats exactly what I was missing! So far I used some list comprehension and scipy, which is... not ideal. For completeness sake and if someone needs it in the future, here my code:

    def apply_augmentation(pointclouds:torch.Tensor, rotations:torch.Tensor) -> torch.Tensor:
        """
        Apply augmentation to a batch of point clouds.

        Args:
            pointclouds (torch.Tensor): Input tensor representing a batch of point clouds.
                Shape: (batch_size, num_points, point_dim).
            rotations (torch.Tensor): Input tensor representing rotation angles for each point cloud.
                Shape: (batch_size, 3), where each row specifies rotation angles around X, Y, and Z axes.

        Returns:
            torch.Tensor: Augmented point clouds after rotation and jitter.
                Shape: Same as input 'pointclouds'.
        """
        pointclouds = pointclouds.to(args.device)
        center = torch.mean(pointclouds, dim=1)
        rotation_matrices = euler_angles_to_matrix(rotations, "XYZ").to(args.device)
        centered_pointcloud = pointclouds - center[:, None, :]
        rotated_pointcloud = torch.bmm(centered_pointcloud, rotation_matrices)

        noise = torch.randn(rotated_pointcloud.size()).to(args.device) * args.jitter_percentage
        rotated_pointcloud += noise

        return rotated_pointcloud