Pointcloud normals estimation.

Summary: Estimates normals of a point cloud.

Reviewed By: gkioxari

Differential Revision: D20860182

fbshipit-source-id: 652ec2743fa645e02c01ffa37c2971bf27b89cef

tests/test_ply_io.py

@@ -185,6 +185,18 @@ class TestMeshPlyIO(TestCaseMixin, unittest.TestCase):
         self.assertClose(expected_verts, actual_verts)
         self.assertClose(expected_faces, actual_faces)
 
+    def test_normals_save(self):
+        verts = torch.tensor(
+            [[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=torch.float32
+        )
+        faces = torch.tensor([[0, 1, 2], [0, 2, 3]])
+        normals = torch.tensor(
+            [[0, 1, 0], [1, 0, 0], [0, 0, 1], [1, 0, 0]], dtype=torch.float32
+        )
+        file = StringIO()
+        save_ply(file, verts=verts, faces=faces, verts_normals=normals)
+        file.close()
+
     def test_empty_save_load(self):
         # Vertices + empty faces
         verts = torch.tensor([[0.1, 0.2, 0.3]])

tests/test_points_normals.py

@@ -0,0 +1,152 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+import unittest
+from typing import Tuple, Union
+
+import torch
+from common_testing import TestCaseMixin
+from pytorch3d.ops import (
+    estimate_pointcloud_local_coord_frames,
+    estimate_pointcloud_normals,
+)
+from pytorch3d.structures.pointclouds import Pointclouds
+
+
+DEBUG = False
+
+
+class TestPCLNormals(TestCaseMixin, unittest.TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        torch.manual_seed(42)
+
+    @staticmethod
+    def init_spherical_pcl(
+        batch_size=3, num_points=3000, device=None, use_pointclouds=False
+    ) -> Tuple[Union[torch.Tensor, Pointclouds], torch.Tensor]:
+        # random spherical point cloud
+        pcl = torch.randn(
+            (batch_size, num_points, 3), device=device, dtype=torch.float32
+        )
+        pcl = torch.nn.functional.normalize(pcl, dim=2)
+
+        # GT normals are the same as
+        # the location of each point on the 0-centered sphere
+        normals = pcl.clone()
+
+        # scale and offset the sphere randomly
+        pcl *= torch.rand(batch_size, 1, 1).type_as(pcl) + 1.0
+        pcl += torch.randn(batch_size, 1, 3).type_as(pcl)
+
+        if use_pointclouds:
+            num_points = torch.randint(
+                size=(batch_size,), low=int(num_points * 0.7), high=num_points
+            )
+            pcl, normals = [
+                [x[:np] for x, np in zip(X, num_points)] for X in (pcl, normals)
+            ]
+            pcl = Pointclouds(pcl, normals=normals)
+
+        return pcl, normals
+
+    def test_pcl_normals(self, batch_size=3, num_points=300, neighborhood_size=50):
+        """
+        Tests the normal estimation on a spherical point cloud, where
+        we know the ground truth normals.
+        """
+        device = torch.device("cuda:0")
+        # run several times for different random point clouds
+        for run_idx in range(3):
+            # either use tensors or Pointclouds as input
+            for use_pointclouds in (True, False):
+                # get a spherical point cloud
+                pcl, normals_gt = TestPCLNormals.init_spherical_pcl(
+                    num_points=num_points,
+                    batch_size=batch_size,
+                    device=device,
+                    use_pointclouds=use_pointclouds,
+                )
+                if use_pointclouds:
+                    normals_gt = pcl.normals_padded()
+                    num_pcl_points = pcl.num_points_per_cloud()
+                else:
+                    num_pcl_points = [pcl.shape[1]] * batch_size
+
+                # check for both disambiguation options
+                for disambiguate_directions in (True, False):
+                    curvatures, local_coord_frames = estimate_pointcloud_local_coord_frames(
+                        pcl,
+                        neighborhood_size=neighborhood_size,
+                        disambiguate_directions=disambiguate_directions,
+                    )
+
+                    # estimate the normals
+                    normals = estimate_pointcloud_normals(
+                        pcl,
+                        neighborhood_size=neighborhood_size,
+                        disambiguate_directions=disambiguate_directions,
+                    )
+
+                    # TODO: temporarily disabled
+                    if use_pointclouds:
+                        # test that the class method gives the same output
+                        normals_pcl = pcl.estimate_normals(
+                            neighborhood_size=neighborhood_size,
+                            disambiguate_directions=disambiguate_directions,
+                            assign_to_self=True,
+                        )
+                        normals_from_pcl = pcl.normals_padded()
+                        for nrm, nrm_from_pcl, nrm_pcl, np in zip(
+                            normals, normals_from_pcl, normals_pcl, num_pcl_points
+                        ):
+                            self.assertClose(nrm[:np], nrm_pcl[:np], atol=1e-5)
+                            self.assertClose(nrm[:np], nrm_from_pcl[:np], atol=1e-5)
+
+                    # check that local coord frames give the same normal
+                    # as normals
+                    for nrm, lcoord, np in zip(
+                        normals, local_coord_frames, num_pcl_points
+                    ):
+                        self.assertClose(nrm[:np], lcoord[:np, :, 0], atol=1e-5)
+
+                    # dotp between normals and normals_gt
+                    normal_parallel = (normals_gt * normals).sum(2)
+
+                    # check that normals are on average
+                    # parallel to the expected ones
+                    for normp, np in zip(normal_parallel, num_pcl_points):
+                        abs_parallel = normp[:np].abs()
+                        avg_parallel = abs_parallel.mean()
+                        std_parallel = abs_parallel.std()
+                        self.assertClose(
+                            avg_parallel, torch.ones_like(avg_parallel), atol=1e-2
+                        )
+                        self.assertClose(
+                            std_parallel, torch.zeros_like(std_parallel), atol=1e-2
+                        )
+
+                    if disambiguate_directions:
+                        # check that 95% of normal dot products
+                        # have the same sign
+                        for normp, np in zip(normal_parallel, num_pcl_points):
+                            n_pos = (normp[:np] > 0).sum()
+                            self.assertTrue((n_pos > np * 0.95) or (n_pos < np * 0.05))
+
+                    if DEBUG and run_idx == 0 and not use_pointclouds:
+                        import os
+                        from pytorch3d.io.ply_io import save_ply
+
+                        # export to .ply
+                        outdir = "/tmp/pt3d_pcl_normals_test/"
+                        os.makedirs(outdir, exist_ok=True)
+                        plyfile = os.path.join(
+                            outdir, f"pcl_disamb={disambiguate_directions}.ply"
+                        )
+                        print(f"Storing point cloud with normals to {plyfile}.")
+                        pcl_idx = 0
+                        save_ply(
+                            plyfile,
+                            pcl[pcl_idx].cpu(),
+                            faces=None,
+                            verts_normals=normals[pcl_idx].cpu(),
+                        )