fix self assign for normals est

Summary: Fix self assignment of normals when estimating normals

Reviewed By: nikhilaravi

Differential Revision: D21315980

fbshipit-source-id: 2aa5864c3f066e39e07343f192cc6423ce1ae771

pytorch3d/structures/pointclouds.py

@@ -382,11 +382,9 @@ class Pointclouds(object):
             if self._normals_padded is None:
                 # No normals provided so return None
                 return None
-            self._normals_list = []
-            for i in range(self._N):
-                self._normals_list.append(
-                    self._normals_padded[i, : self.num_points_per_cloud()[i]]
-                )
+            self._normals_list = struct_utils.padded_to_list(
+                self._normals_padded, self.num_points_per_cloud().tolist()
+            )
         return self._normals_list
 
     def features_list(self):
@@ -400,11 +398,9 @@ class Pointclouds(object):
             if self._features_padded is None:
                 # No features provided so return None
                 return None
-            self._features_list = []
-            for i in range(self._N):
-                self._features_list.append(
-                    self._features_padded[i, : self.num_points_per_cloud()[i]]
-                )
+            self._features_list = struct_utils.padded_to_list(
+                self._features_padded, self.num_points_per_cloud().tolist()
+            )
         return self._features_list
 
     def points_packed(self):
@@ -887,9 +883,14 @@ class Pointclouds(object):
 
         # assign to self
         if assign_to_self:
-            self._normals_list, self._normals_padded, _ = self._parse_auxiliary_input(
-                normals_est
-            )
+            _, self._normals_padded, _ = self._parse_auxiliary_input(normals_est)
+            self._normals_list, self._normals_packed = None, None
+            if self._points_list is not None:
+                # update self._normals_list
+                self.normals_list()
+            if self._points_packed is not None:
+                # update self._normals_packed
+                self._normals_packed = torch.cat(self._normals_list, dim=0)
 
         return normals_est
 

tests/test_pointclouds.py

@@ -6,6 +6,7 @@ import unittest
 import numpy as np
 import torch
 from common_testing import TestCaseMixin
+from pytorch3d.structures import utils as struct_utils
 from pytorch3d.structures.pointclouds import Pointclouds
 
 
@@ -22,6 +23,7 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
         lists_to_tensors: bool = False,
         with_normals: bool = True,
         with_features: bool = True,
+        min_points: int = 0,
     ):
         """
         Function to generate a Pointclouds object of N meshes with
@@ -36,12 +38,13 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
                               tensors (=True) of points/normals/features.
             with_normals: bool whether to include normals
             with_features: bool whether to include features
+            min_points: Min number of points per cloud
 
         Returns:
             Pointclouds object.
         """
         device = torch.device("cuda:0")
-        p = torch.randint(max_points, size=(num_clouds,))
+        p = torch.randint(low=min_points, high=max_points, size=(num_clouds,))
         if lists_to_tensors:
             p.fill_(p[0])
 
@@ -893,6 +896,35 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "Input box must be of shape"):
             clouds.inside_box(invalid_box)
 
+    def test_estimate_normals(self):
+        for with_normals in (True, False):
+            for run_padded in (True, False):
+                for run_packed in (True, False):
+
+                    clouds = TestPointclouds.init_cloud(
+                        3,
+                        100,
+                        with_normals=with_normals,
+                        with_features=False,
+                        min_points=60,
+                    )
+                    nums = clouds.num_points_per_cloud()
+                    if run_padded:
+                        clouds.points_padded()
+                    if run_packed:
+                        clouds.points_packed()
+
+                    normals_est_padded = clouds.estimate_normals(assign_to_self=True)
+                    normals_est_list = struct_utils.padded_to_list(
+                        normals_est_padded, nums.tolist()
+                    )
+                    self.assertClose(clouds.normals_padded(), normals_est_padded)
+                    for i in range(len(clouds)):
+                        self.assertClose(clouds.normals_list()[i], normals_est_list[i])
+                    self.assertClose(
+                        clouds.normals_packed(), torch.cat(normals_est_list, dim=0)
+                    )
+
     @staticmethod
     def compute_packed_with_init(
         num_clouds: int = 10, max_p: int = 100, features: int = 300