join_pointclouds_as_scene

Summary: New function

Reviewed By: davidsonic

Differential Revision: D42776590

fbshipit-source-id: 2a6e73480bcf2d1749f86bcb22d1942e3e8d3167

Author: bottler

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
 

pytorch3d/structures/pointclouds.py

@@ -124,12 +124,14 @@ def __init__(self, points, normals=None, features=None) -> None:
             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

tests/test_pointclouds.py

@@ -11,7 +11,11 @@
 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 @@ def check_triple(points, points3):
         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 @@ def check_triple(points, points3):
         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