diff --git a/pytorch3d/implicitron/tools/point_cloud_utils.py b/pytorch3d/implicitron/tools/point_cloud_utils.py index 904e4567..5954aace 100644 --- 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 index f1b5097f..dc3176cc 100644 --- a/tests/implicitron/test_pointcloud_utils.py +++ b/tests/implicitron/test_pointcloud_utils.py @@ -62,3 +62,11 @@ class TestPointCloudUtils(TestCaseMixin, unittest.TestCase): ) [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, + )