From a123815f40d19127ad8dc20ae93483303220046f Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Thu, 9 Mar 2023 06:51:13 -0800 Subject: [PATCH] join_pointclouds_as_scene Summary: New function Reviewed By: davidsonic Differential Revision: D42776590 fbshipit-source-id: 2a6e73480bcf2d1749f86bcb22d1942e3e8d3167 --- pytorch3d/structures/__init__.py | 6 ++++- pytorch3d/structures/pointclouds.py | 38 +++++++++++++++++++++++++++++ tests/test_pointclouds.py | 22 ++++++++++++++--- 3 files changed, 61 insertions(+), 5 deletions(-) diff --git a/pytorch3d/structures/__init__.py b/pytorch3d/structures/__init__.py index 37986ad7..b92e8724 100644 --- a/pytorch3d/structures/__init__.py +++ b/pytorch3d/structures/__init__.py @@ -5,7 +5,11 @@ # LICENSE file in the root directory of this source tree. from .meshes import join_meshes_as_batch, join_meshes_as_scene, Meshes -from .pointclouds import Pointclouds +from .pointclouds import ( + join_pointclouds_as_batch, + join_pointclouds_as_scene, + Pointclouds, +) from .utils import list_to_packed, list_to_padded, packed_to_list, padded_to_list from .volumes import Volumes diff --git a/pytorch3d/structures/pointclouds.py b/pytorch3d/structures/pointclouds.py index bf443988..654b5a26 100644 --- a/pytorch3d/structures/pointclouds.py +++ b/pytorch3d/structures/pointclouds.py @@ -124,12 +124,14 @@ class Pointclouds: normals: Can be either + - None - List where each element is a tensor of shape (num_points, 3) containing the normal vector for each point. - Padded float tensor of shape (num_clouds, num_points, 3). features: Can be either + - None - List where each element is a tensor of shape (num_points, C) containing the features for the points in the cloud. - Padded float tensor of shape (num_clouds, num_points, C). @@ -1260,6 +1262,42 @@ def join_pointclouds_as_batch(pointclouds: Sequence[Pointclouds]) -> Pointclouds field_list = None else: field_list = [p for points in field_list for p in points] + if field == "features" and any( + p.shape[1] != field_list[0].shape[1] for p in field_list[1:] + ): + raise ValueError("Pointclouds must have the same number of features") kwargs[field] = field_list return Pointclouds(**kwargs) + + +def join_pointclouds_as_scene( + pointclouds: Union[Pointclouds, List[Pointclouds]] +) -> Pointclouds: + """ + Joins a batch of point cloud in the form of a Pointclouds object or a list of Pointclouds + objects as a single point cloud. If the input is a list, the Pointclouds objects in the + list must all be on the same device, and they must either all or none have features and + all or none have normals. + + Args: + Pointclouds: Pointclouds object that contains a batch of point clouds, or a list of + Pointclouds objects. + + Returns: + new Pointclouds object containing a single point cloud + """ + if isinstance(pointclouds, list): + pointclouds = join_pointclouds_as_batch(pointclouds) + + if len(pointclouds) == 1: + return pointclouds + points = pointclouds.points_packed() + features = pointclouds.features_packed() + normals = pointclouds.normals_packed() + pointcloud = Pointclouds( + points=points[None], + features=None if features is None else features[None], + normals=None if normals is None else normals[None], + ) + return pointcloud diff --git a/tests/test_pointclouds.py b/tests/test_pointclouds.py index e8dc3d12..9ad80d50 100644 --- a/tests/test_pointclouds.py +++ b/tests/test_pointclouds.py @@ -11,7 +11,11 @@ import unittest import numpy as np import torch from pytorch3d.structures import utils as struct_utils -from pytorch3d.structures.pointclouds import join_pointclouds_as_batch, Pointclouds +from pytorch3d.structures.pointclouds import ( + join_pointclouds_as_batch, + join_pointclouds_as_scene, + Pointclouds, +) from .common_testing import TestCaseMixin @@ -1159,9 +1163,9 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase): normals = [torch.rand(length, 3) for length in lengths] # Test with normals and features present - pcl = Pointclouds(points=points, features=features, normals=normals) - pcl3 = join_pointclouds_as_batch([pcl] * 3) - check_triple(pcl, pcl3) + pcl1 = Pointclouds(points=points, features=features, normals=normals) + pcl3 = join_pointclouds_as_batch([pcl1] * 3) + check_triple(pcl1, pcl3) # Test with normals and features present for tensor backed pointclouds N, P, D = 5, 30, 4 @@ -1173,15 +1177,25 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase): pcl3 = join_pointclouds_as_batch([pcl] * 3) check_triple(pcl, pcl3) + # Test with inconsistent #features + with self.assertRaisesRegex(ValueError, "same number of features"): + join_pointclouds_as_batch([pcl1, pcl]) + # Test without normals pcl_nonormals = Pointclouds(points=points, features=features) pcl3 = join_pointclouds_as_batch([pcl_nonormals] * 3) check_triple(pcl_nonormals, pcl3) + pcl_scene = join_pointclouds_as_scene([pcl_nonormals] * 3) + self.assertEqual(len(pcl_scene), 1) + self.assertClose(pcl_scene.features_packed(), pcl3.features_packed()) # Test without features pcl_nofeats = Pointclouds(points=points, normals=normals) pcl3 = join_pointclouds_as_batch([pcl_nofeats] * 3) check_triple(pcl_nofeats, pcl3) + pcl_scene = join_pointclouds_as_scene([pcl_nofeats] * 3) + self.assertEqual(len(pcl_scene), 1) + self.assertClose(pcl_scene.normals_packed(), pcl3.normals_packed()) # Check error raised if all pointclouds in the batch # are not consistent in including normals/features