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
This commit is contained in:
Ido Zachevsky 2024-01-16 03:38:26 -08:00 committed by Facebook GitHub Bot
parent 292acc71a3
commit 799c1cd21b
2 changed files with 11 additions and 1 deletions

View File

@ -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])

View File

@ -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,
)