doc rgbd point cloud

Summary: docstring and shape fix

Reviewed By: shapovalov

Differential Revision: D42609661

fbshipit-source-id: fd50234872ad61b5452821eeb89d51344f70c957

tests/implicitron/test_pointcloud_utils.py

@@ -0,0 +1,64 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from pytorch3d.implicitron.tools.point_cloud_utils import get_rgbd_point_cloud
+
+from pytorch3d.renderer.cameras import PerspectiveCameras
+from tests.common_testing import TestCaseMixin
+
+
+class TestPointCloudUtils(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+
+    def test_unproject(self):
+        H, W = 50, 100
+
+        # Random RGBD image with depth 3
+        # (depth 0 = at the camera)
+        # and purple in the upper right corner
+
+        image = torch.rand(4, H, W)
+        depth = 3
+        image[3] = depth
+        image[1, H // 2 :, W // 2 :] *= 0.4
+
+        # two ways to define the same camera:
+        # at the origin facing the positive z axis
+        ndc_camera = PerspectiveCameras(focal_length=1.0)
+        screen_camera = PerspectiveCameras(
+            focal_length=H // 2,
+            in_ndc=False,
+            image_size=((H, W),),
+            principal_point=((W / 2, H / 2),),
+        )
+
+        for camera in (ndc_camera, screen_camera):
+            # 1. z-depth
+            cloud = get_rgbd_point_cloud(
+                camera,
+                image_rgb=image[:3][None],
+                depth_map=image[3:][None],
+                euclidean=False,
+            )
+            [points] = cloud.points_list()
+            self.assertConstant(points[:, 2], depth)  # constant depth
+            extremes = depth * torch.tensor([W / H - 1 / H, 1 - 1 / H])
+            self.assertClose(points[:, :2].min(0).values, -extremes)
+            self.assertClose(points[:, :2].max(0).values, extremes)
+
+            # 2. euclidean
+            cloud = get_rgbd_point_cloud(
+                camera,
+                image_rgb=image[:3][None],
+                depth_map=image[3:][None],
+                euclidean=True,
+            )
+            [points] = cloud.points_list()
+            self.assertConstant(torch.norm(points, dim=1), depth, atol=1e-5)