diff --git a/pytorch3d/structures/pointclouds.py b/pytorch3d/structures/pointclouds.py index 6f6fc223..082507ca 100644 --- a/pytorch3d/structures/pointclouds.py +++ b/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 diff --git a/tests/test_pointclouds.py b/tests/test_pointclouds.py index fe877791..efc02a2b 100644 --- a/tests/test_pointclouds.py +++ b/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