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
facebookresearch · Jan 16, 2024 · 799c1cd · 799c1cd
1 parent 292acc7
commit 799c1cd
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/pytorch3d/implicitron/tools/point_cloud_utils.py b/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])
 

diff --git a/tests/implicitron/test_pointcloud_utils.py b/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,
+            )