From 799c1cd21beff84e50ac4ab7a480e715780da2de Mon Sep 17 00:00:00 2001 From: Ido Zachevsky Date: Tue, 16 Jan 2024 03:38:26 -0800 Subject: [PATCH] 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 --- pytorch3d/implicitron/tools/point_cloud_utils.py | 4 +++- tests/implicitron/test_pointcloud_utils.py | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/pytorch3d/implicitron/tools/point_cloud_utils.py b/pytorch3d/implicitron/tools/point_cloud_utils.py index 904e45675..5954aace6 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 f1b5097fd..dc3176cc6 100644 --- 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, + )