subsample pointclouds

Summary: New function to randomly subsample Pointclouds to a maximum size.

Reviewed By: nikhilaravi

Differential Revision: D30936533

fbshipit-source-id: 789eb5004b6a233034ec1c500f20f2d507a303ff
This commit is contained in:
Jeremy Reizenstein
2021-10-02 13:37:09 -07:00
committed by Facebook GitHub Bot
parent ee2b2feb98
commit 4281df19ce
3 changed files with 92 additions and 21 deletions

View File

@@ -1057,6 +1057,45 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
clouds.normals_packed(), torch.cat(normals_est_list, dim=0)
)
def test_subsample(self):
lengths = [4, 5, 13, 3]
points = [torch.rand(length, 3) for length in lengths]
features = [torch.rand(length, 5) for length in lengths]
normals = [torch.rand(length, 3) for length in lengths]
pcl1 = Pointclouds(points=points).cuda()
self.assertIs(pcl1, pcl1.subsample(13))
self.assertIs(pcl1, pcl1.subsample([6, 13, 13, 13]))
lengths_max_4 = torch.tensor([4, 4, 4, 3]).cuda()
for with_normals, with_features in itertools.product([True, False], repeat=2):
with self.subTest(f"{with_normals} {with_features}"):
pcl = Pointclouds(
points=points,
normals=normals if with_normals else None,
features=features if with_features else None,
)
pcl_copy = pcl.subsample(max_points=4)
for length, points_ in zip(lengths_max_4, pcl_copy.points_list()):
self.assertEqual(points_.shape, (length, 3))
if with_normals:
for length, normals_ in zip(lengths_max_4, pcl_copy.normals_list()):
self.assertEqual(normals_.shape, (length, 3))
else:
self.assertIsNone(pcl_copy.normals_list())
if with_features:
for length, features_ in zip(
lengths_max_4, pcl_copy.features_list()
):
self.assertEqual(features_.shape, (length, 5))
else:
self.assertIsNone(pcl_copy.features_list())
pcl2 = Pointclouds(points=points)
pcl_copy2 = pcl2.subsample(lengths_max_4)
for length, points_ in zip(lengths_max_4, pcl_copy2.points_list()):
self.assertEqual(points_.shape, (length, 3))
@staticmethod
def compute_packed_with_init(
num_clouds: int = 10, max_p: int = 100, features: int = 300