[perception] Add PointCloud::Crop and Concatenate(vector<PointCloud>) (…

…RobotLocomotion#17829)

* [perception] Add PointCloud::Crop and Concatenate(vector<PointCloud>)

Towards RussTedrake/manipulation#130.

bindings/pydrake/perception_py.cc

@@ -118,11 +118,15 @@ void init_perception(py::module m) {
             [](PointCloud* self, const PointCloud& other) {
               self->SetFrom(other);
             },
-            py::arg("other"), cls_doc.SetFrom.doc);
+            py::arg("other"), cls_doc.SetFrom.doc)
+        .def("Crop", &Class::Crop, py::arg("lower_xyz"), py::arg("upper_xyz"),
+            cls_doc.Crop.doc);
   }
 
   AddValueInstantiation<PointCloud>(m);
 
+  m.def("Concatenate", &Concatenate, py::arg("clouds"), doc.Concatenate.doc);
+
   {
     using Class = DepthImageToPointCloud;
     constexpr auto& cls_doc = doc.DepthImageToPointCloud;

bindings/pydrake/test/perception_test.py

@@ -82,6 +82,15 @@ def test_point_cloud_api(self):
         # Test Systems' value registration.
         self.assertIsInstance(AbstractValue.Make(pc), Value[mut.PointCloud])
 
+        pc = mut.PointCloud(new_size=2, fields=mut.Fields(mut.BaseField.kXYZs))
+        test_xyzs = [[1., 2., 3.], [4., 5., 6.]]
+        pc.mutable_xyzs().T[:] = test_xyzs
+        crop = pc.Crop(lower_xyz=[3, 4, 5], upper_xyz=[5, 6, 7])
+        self.assertEqual(crop.size(), 1)
+
+        pc_merged = mut.Concatenate(clouds=[pc, pc_new])
+        self.assertEqual(pc_merged.size(), pc.size() + pc_new.size())
+
     def test_depth_image_to_point_cloud_api(self):
         camera_info = CameraInfo(width=640, height=480, fov_y=np.pi / 4)
         dut = mut.DepthImageToPointCloud(camera_info=camera_info)

perception/point_cloud.cc

@@ -6,6 +6,7 @@
 #include <fmt/ostream.h>
 
 #include "drake/common/drake_assert.h"
+#include "drake/common/drake_throw.h"
 
 using Eigen::Map;
 using Eigen::NoChange;
@@ -318,5 +319,64 @@ void PointCloud::RequireExactFields(
   }
 }
 
+PointCloud PointCloud::Crop(const Eigen::Ref<const Vector3<T>>& lower_xyz,
+                            const Eigen::Ref<const Vector3<T>>& upper_xyz) {
+  DRAKE_DEMAND((lower_xyz.array() <= upper_xyz.array()).all());
+  if (!has_xyzs()) {
+    throw std::runtime_error("PointCloud must have xyzs in order to Crop");
+  }
+  PointCloud crop(size_, fields(), true);
+  int index = 0;
+  for (int i = 0; i < size_; ++i) {
+    if (((xyzs().col(i).array() >= lower_xyz.array()) &&
+         (xyzs().col(i).array() <= upper_xyz.array()))
+            .all()) {
+      crop.mutable_xyzs().col(index) = xyzs().col(i);
+      if (has_normals()) {
+        crop.mutable_normals().col(index) = normals().col(i);
+      }
+      if (has_rgbs()) {
+        crop.mutable_rgbs().col(index) = rgbs().col(i);
+      }
+      if (has_descriptors()) {
+        crop.mutable_descriptors().col(index) = descriptors().col(i);
+      }
+      ++index;
+    }
+  }
+  crop.resize(index);
+  return crop;
+}
+
+PointCloud Concatenate(const std::vector<PointCloud>& clouds) {
+  const int num_clouds = clouds.size();
+  DRAKE_DEMAND(num_clouds >= 1);
+  int count = clouds[0].size();
+  for (int i = 1; i < num_clouds; ++i) {
+    DRAKE_THROW_UNLESS(clouds[i].fields() == clouds[0].fields());
+    count += clouds[i].size();
+  }
+  PointCloud new_cloud(count, clouds[0].fields(), true);
+  int index = 0;
+  for (int i = 0; i < num_clouds; ++i) {
+    const int s = clouds[i].size();
+    if (new_cloud.has_normals()) {
+      new_cloud.mutable_xyzs().middleCols(index, s) = clouds[i].xyzs();
+    }
+    if (new_cloud.has_normals()) {
+      new_cloud.mutable_normals().middleCols(index, s) = clouds[i].normals();
+    }
+    if (new_cloud.has_rgbs()) {
+      new_cloud.mutable_rgbs().middleCols(index, s) = clouds[i].rgbs();
+    }
+    if (new_cloud.has_descriptors()) {
+      new_cloud.mutable_descriptors().middleCols(index, s) =
+          clouds[i].descriptors();
+    }
+    index += s;
+  }
+  return new_cloud;
+}
+
 }  // namespace perception
 }  // namespace drake

perception/point_cloud.h

@@ -300,6 +300,19 @@ class PointCloud final {
 
   /// @}
 
+  /// @name Point Cloud Processing
+  /// @{
+
+  /// Returns a new point cloud containing only the points in `this` with xyz
+  /// values within the axis-aligned bounding box defined by `lower_xyz` and
+  /// `upper_xyz`. Requires that xyz values are defined.
+  /// @pre lower_xyz <= upper_xyz (elementwise).
+  /// @throws std::exception if has_xyzs() != true.
+  PointCloud Crop(const Eigen::Ref<const Vector3<T>>& lower_xyz,
+            const Eigen::Ref<const Vector3<T>>& upper_xyz);
+
+  /// @}
+
   // TODO(eric.cousineau): Add storage for indices, with SHOT as a motivating
   // example.
 
@@ -320,6 +333,12 @@ class PointCloud final {
   std::unique_ptr<Storage> storage_;
 };
 
+/// Returns a new point cloud that includes all of the points from the point
+/// clouds in `clouds`. All of the `clouds` must have the same fields.
+/// @pre `clouds` contains at least one point cloud.
+/// @throws std::exception if the clouds have different fields defined.
+PointCloud Concatenate(const std::vector<PointCloud>& clouds);
+
 // TODO(eric.cousineau): Consider a way of reinterpret_cast<>ing the array
 // data to permit more semantic access to members, PCL-style
 // (e.g. point.x, point.r, etc: the reverse of PointCloud<>::getMatrixXfMap()).

perception/test/point_cloud_test.cc

@@ -196,6 +196,33 @@ GTEST_TEST(PointCloudTest, Basic) {
                 return cloud.mutable_descriptor(i);
               },
               [](PointCloud& cloud, int i) { return cloud.descriptor(i); });
+
+  {  // Crop
+    PointCloud cloud(count, pc_flags::kXYZs | pc_flags::kNormals |
+                                pc_flags::kRGBs |
+                                pc_flags::kDescriptorCurvature);
+    cloud.mutable_xyzs() = xyzs_expected;
+    cloud.mutable_rgbs() = rgbs_expected;
+    cloud.mutable_normals() = normals_expected;
+    cloud.mutable_descriptors() = descriptors_expected;
+
+    PointCloud cropped =
+        cloud.Crop(Eigen::Vector3f{4, 5, 6}, Eigen::Vector3f{10, 20, 30});
+    EXPECT_EQ(cropped.size(), 2);
+    std::vector<int> indices{1, 3};
+
+    for (int i = 0; i < static_cast<int>(indices.size()); ++i) {
+      EXPECT_TRUE(CompareMatrices(cropped.xyzs().col(i),
+                                  xyzs_expected.col(indices[i])));
+      EXPECT_TRUE((cropped.rgbs().col(i).array() ==
+                   rgbs_expected.col(indices[i]).array())
+                      .all());
+      EXPECT_TRUE(CompareMatrices(cropped.normals().col(i),
+                                  normals_expected.col(indices[i])));
+      EXPECT_TRUE(CompareMatrices(cropped.descriptors().col(i),
+                                  descriptors_expected.col(indices[i])));
+    }
+  }
 }
 
 GTEST_TEST(PointCloudTest, Fields) {
@@ -257,6 +284,53 @@ GTEST_TEST(PointCloudTest, Fields) {
   }
 }
 
+GTEST_TEST(PointCloud, Concatenate) {
+  auto fields = pc_flags::kXYZs | pc_flags::kNormals | pc_flags::kRGBs |
+                pc_flags::kDescriptorCurvature;
+
+  std::vector<PointCloud> clouds;
+  // Populate three point clouds with arbitrary values.
+  clouds.push_back(PointCloud(2, fields, true));
+  clouds.push_back(PointCloud(3, fields, true));
+  clouds.push_back(PointCloud(1, fields, true));
+
+  clouds[0].mutable_xyzs() << 1, 2, 3, 4, 5, 6;
+  clouds[0].mutable_rgbs() << 10, 20, 30, 40, 50, 60;
+  clouds[0].mutable_normals() << 3, 4, -1, 12, 4, 32;
+  clouds[0].mutable_descriptors() << 8, 2;
+
+  clouds[1].mutable_xyzs() << 421, 1, 1, 4, 4, 2, 141, 1, 63;
+  clouds[1].mutable_rgbs() << 57, 1, 42, 25, 25, 10, 12, 63, 192;
+  clouds[1].mutable_normals() << -2, 1, 532, 7, 2, 6, 36, 84, 42;
+  clouds[1].mutable_descriptors() << 10, 52, 1;
+
+  clouds[2].mutable_xyzs() << 71, 2, 1;
+  clouds[2].mutable_rgbs() << 53, 24, 172;
+  clouds[2].mutable_normals() << 91, 4, 44;
+  clouds[2].mutable_descriptors() << 91;
+
+  PointCloud merged = Concatenate(clouds);
+  EXPECT_EQ(merged.size(), 6);
+
+  Eigen::Matrix3Xf xyzs_expected(3, 6);
+  xyzs_expected << clouds[0].xyzs(), clouds[1].xyzs(), clouds[2].xyzs();
+  EXPECT_TRUE(CompareMatrices(merged.xyzs(), xyzs_expected));
+
+  Matrix3X<uint8_t> rgbs_expected(3, 6);
+  rgbs_expected << clouds[0].rgbs(), clouds[1].rgbs(), clouds[2].rgbs();
+  EXPECT_TRUE((merged.rgbs().array() == rgbs_expected.array()).all());
+
+  Eigen::Matrix3Xf normals_expected(3, 6);
+  normals_expected << clouds[0].normals(), clouds[1].normals(),
+      clouds[2].normals();
+  EXPECT_TRUE(CompareMatrices(merged.normals(), normals_expected));
+
+  Eigen::RowVectorXf descriptors_expected(6);
+  descriptors_expected << clouds[0].descriptors(), clouds[1].descriptors(),
+      clouds[2].descriptors();
+  EXPECT_TRUE(CompareMatrices(merged.descriptors(), descriptors_expected));
+}
+
 }  // namespace
 }  // namespace perception
 }  // namespace drake