Allow get_rgbd_point_cloud to take any #channels

Summary: Fixed `get_rgbd_point_cloud` to take any number of image input channels.

Reviewed By: bottler

Differential Revision: D52796276

fbshipit-source-id: 3ddc0d1e337a6cc53fc86c40a6ddb136f036f9bc

Author: Ido Zachevsky

pytorch3d/implicitron/tools/point_cloud_utils.py

@@ -80,7 +80,9 @@ def get_rgbd_point_cloud(
         mode="bilinear",
         align_corners=False,
     )
-    pts_colors = pts_colors.permute(0, 2, 3, 1).reshape(-1, 3)[pts_mask]
+    pts_colors = pts_colors.permute(0, 2, 3, 1).reshape(-1, image_rgb.shape[1])[
+        pts_mask
+    ]
 
     return Pointclouds(points=pts_3d[None], features=pts_colors[None])
 

tests/implicitron/test_pointcloud_utils.py

@@ -62,3 +62,11 @@ def test_unproject(self):
             )
             [points] = cloud.points_list()
             self.assertConstant(torch.norm(points, dim=1), depth, atol=1e-5)
+
+            # 3. four channels
+            get_rgbd_point_cloud(
+                camera,
+                image_rgb=image[None],
+                depth_map=image[3:][None],
+                euclidean=True,
+            )