[perception] Implement PointCloud::EstimateNormals.

Adds nanoflann as a (header-only) dependency, and implements the basic
normal estimation algorithm for PointCloud, as seen in PCL and/or
Open3d.

Towards RussTedrake/manipulation#130.

bindings/pydrake/perception_py.cc

@@ -122,7 +122,9 @@ void init_perception(py::module m) {
         .def("Crop", &Class::Crop, py::arg("lower_xyz"), py::arg("upper_xyz"),
             cls_doc.Crop.doc)
         .def("VoxelizedDownSample", &Class::VoxelizedDownSample,
-            py::arg("voxel_size"), cls_doc.VoxelizedDownSample.doc);
+            py::arg("voxel_size"), cls_doc.VoxelizedDownSample.doc)
+        .def("EstimateNormals", &Class::EstimateNormals, py::arg("distance"),
+            py::arg("max_nearest_neighbors"), cls_doc.EstimateNormals.doc);
   }
 
   AddValueInstantiation<PointCloud>(m);

bindings/pydrake/test/perception_test.py

@@ -94,6 +94,10 @@ def test_point_cloud_api(self):
         pc_downsampled = pc_merged.VoxelizedDownSample(voxel_size=2.0)
         self.assertIsInstance(pc_downsampled, mut.PointCloud)
 
+        self.assertFalse(pc_merged.has_normals())
+        pc_merged.EstimateNormals(distance=1, max_nearest_neighbors=50)
+        self.assertTrue(pc_merged.has_normals())
+
     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/BUILD.bazel

@@ -42,6 +42,7 @@ drake_cc_library(
     deps = [
         "//common:unused",
         "@abseil_cpp_internal//absl/container:flat_hash_map",
+        "@nanoflann",
     ],
 )
 
@@ -91,6 +92,7 @@ drake_cc_googletest(
     name = "point_cloud_test",
     deps = [
         ":point_cloud",
+        "//common:random",
         "//common/test_utilities:eigen_matrix_compare",
         "//common/test_utilities:expect_no_throw",
     ],

perception/point_cloud.cc

@@ -8,6 +8,7 @@
 #include "absl/container/flat_hash_map.h"
 #include <fmt/format.h>
 #include <fmt/ostream.h>
+#include <nanoflann.hpp>
 
 #include "drake/common/drake_assert.h"
 #include "drake/common/drake_throw.h"
@@ -62,6 +63,17 @@ class PointCloud::Storage {
     CheckInvariants();
   }
 
+  // Update fields, allocating (but not initializing) new fields when needed.
+  void UpdateFields(pc_flags::Fields f) {
+    xyzs_.conservativeResize(NoChange, f.contains(pc_flags::kXYZs) ? size_ : 0);
+    normals_.conservativeResize(NoChange,
+                                f.contains(pc_flags::kNormals) ? size_ : 0);
+    rgbs_.conservativeResize(NoChange, f.contains(pc_flags::kRGBs) ? size_ : 0);
+    descriptors_.conservativeResize(NoChange, f.has_descriptor() ? size_ : 0);
+    fields_ = f;
+    CheckInvariants();
+  }
+
   Eigen::Ref<Matrix3X<T>> xyzs() { return xyzs_; }
   Eigen::Ref<Matrix3X<T>> normals() { return normals_; }
   Eigen::Ref<Matrix3X<C>> rgbs() { return rgbs_; }
@@ -87,7 +99,7 @@ class PointCloud::Storage {
     }
   }
 
-  const pc_flags::Fields fields_;
+  pc_flags::Fields fields_;
   int size_{};
   Matrix3X<T> xyzs_;
   Matrix3X<T> normals_;
@@ -480,5 +492,81 @@ PointCloud PointCloud::VoxelizedDownSample(double voxel_size) const {
   return down_sampled;
 }
 
+void PointCloud::EstimateNormals(double radius, int max_nearest_neighbors) {
+  DRAKE_DEMAND(radius > 0);
+  DRAKE_DEMAND(max_nearest_neighbors > 0);
+  DRAKE_THROW_UNLESS(size() >= 3);
+  DRAKE_THROW_UNLESS(has_xyzs());
+
+  if (!has_normals()) {
+    fields_ |= pc_flags::kNormals;
+    storage_->UpdateFields(fields_);
+  }
+
+  const Eigen::MatrixX3f data = xyzs().transpose();
+  nanoflann::KDTreeEigenMatrixAdaptor<Eigen::MatrixX3f, 3> kd_tree(3, data);
+
+  // Iterate through all points and compute their normals.
+  // TODO(russt): Make an OpenMP implementation of this.
+  VectorX<int64_t> indices(max_nearest_neighbors);
+  Eigen::VectorXf distances(max_nearest_neighbors);
+  Eigen::Matrix3d covariance;
+  Eigen::Matrix<double, 9, 1> cumulants;
+
+  for (int i = 0; i < size_; ++i) {
+    // With nanoflann, I can either search for the max_nearest_neighbors and
+    // take the ones closer than radius, or search for all points closer than
+    // radius and keep the max_nearest_neighbors.
+    const int num_neighbors = kd_tree.index->knnSearch(
+        xyz(i).data(), max_nearest_neighbors, indices.data(), distances.data());
+
+    if (num_neighbors < 2) {
+      mutable_normal(i) = Eigen::Vector3f{0, 0, 1};
+      continue;
+    }
+
+    // Compute the covariance matrix.
+    {
+      cumulants.setZero();
+      int count = 0;
+      for (int j = 0; j < num_neighbors; ++j) {
+        if (distances[j] <= radius) {
+          ++count;
+          const Eigen::Vector3f& point = xyz(indices[j]);
+          cumulants(0) += point(0);
+          cumulants(1) += point(1);
+          cumulants(2) += point(2);
+          cumulants(3) += point(0) * point(0);
+          cumulants(4) += point(0) * point(1);
+          cumulants(5) += point(0) * point(2);
+          cumulants(6) += point(1) * point(1);
+          cumulants(7) += point(1) * point(2);
+          cumulants(8) += point(2) * point(2);
+        }
+      }
+      if (count < 2) {
+        mutable_normal(i) = Eigen::Vector3f{0, 0, 1};
+        continue;
+      }
+      cumulants /= count;
+      covariance(0, 0) = cumulants(3) - cumulants(0) * cumulants(0);
+      covariance(1, 1) = cumulants(6) - cumulants(1) * cumulants(1);
+      covariance(2, 2) = cumulants(8) - cumulants(2) * cumulants(2);
+      covariance(0, 1) = cumulants(4) - cumulants(0) * cumulants(1);
+      covariance(1, 0) = covariance(0, 1);
+      covariance(0, 2) = cumulants(5) - cumulants(0) * cumulants(2);
+      covariance(2, 0) = covariance(0, 2);
+      covariance(1, 2) = cumulants(7) - cumulants(1) * cumulants(2);
+      covariance(2, 1) = covariance(1, 2);
+    }
+
+    // TODO(russt): Open3d implements a "FastEigen3x3" for an optimized version
+    // of this. We probably should, too.
+    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> solver;
+    solver.compute(covariance, Eigen::ComputeEigenvectors);
+    mutable_normal(i) = solver.eigenvectors().col(0).cast<float>();
+  }
+}
+
 }  // namespace perception
 }  // namespace drake

perception/point_cloud.h

@@ -330,6 +330,14 @@ class PointCloud final {
   /// @throws std::exception if voxel_size <= 0.
   PointCloud VoxelizedDownSample(double voxel_size) const;
 
+  /// Estimates the normal vectors in `this` by fitting a plane at each point
+  /// in the cloud using up to `max_nearest_neighbors` within Euclidean
+  /// distance `radius` from the point. If has_normals() is false, then new
+  /// normals will be allocated (and has_normals() will become true).
+  ///
+  /// @throws std::exception if has_xyzs() is false or if size() < 3.
+  void EstimateNormals(double radius, int max_nearest_neighbors);
+
  private:
   void SetDefault(int start, int num);
 
@@ -339,7 +347,7 @@ class PointCloud final {
   // Represents the size of the point cloud.
   int size_{};
   // Represents which fields are enabled for this point cloud.
-  const pc_flags::Fields fields_{pc_flags::kXYZs};
+  pc_flags::Fields fields_{pc_flags::kXYZs};
   // Owns storage used for the point cloud.
   std::unique_ptr<Storage> storage_;
 };

perception/test/point_cloud_test.cc

@@ -5,11 +5,13 @@
 
 #include <gtest/gtest.h>
 
+#include "drake/common/random.h"
 #include "drake/common/test_utilities/eigen_matrix_compare.h"
 #include "drake/common/test_utilities/expect_no_throw.h"
 
 using Eigen::Matrix3Xf;
 using Eigen::Matrix4Xf;
+using Eigen::Vector3f;
 
 using ::testing::AssertionResult;
 using ::testing::AssertionSuccess;
@@ -381,10 +383,7 @@ GTEST_TEST(PointCloudTest, VoxelizedDownSample) {
         double kTol = 1e-8;
         bool found_match = false;
         int i = 0;
-        std::cout << xyz.transpose() << std::endl;
         for (; i < down_sampled.size(); ++i) {
-          std::cout << i << ": " << down_sampled.xyz(i).transpose()
-                    << std::endl;
           if (CompareMatrices(down_sampled.xyz(i), xyz.cast<float>(), kTol)) {
             found_match = true;
             break;
@@ -412,6 +411,72 @@ GTEST_TEST(PointCloudTest, VoxelizedDownSample) {
   CheckHasPointAveragedFrom({5});
 }
 
+// Checks that normal has unit magnitude and that normal == expected up to a
+// sign flip.
+void CheckNormal(const Eigen::Ref<const Eigen::Vector3f>& normal,
+                 const Eigen::Ref<const Eigen::Vector3f>& expected,
+                 double kTolerance) {
+  EXPECT_NEAR(normal.norm(), 1.0, kTolerance)
+      << "Normal " << normal << " does not have unit length.";
+  EXPECT_NEAR(std::abs(normal.dot(expected)), 1.0, kTolerance)
+      << "normal.dot(expected) = " << normal.dot(expected);
+}
+
+// Test that point clouds that consistent of three points (defining a plane)
+// get estimated normals that match their closed-form solution.
+GTEST_TEST(PointCloudTest, EstimateNormalsPlane) {
+  PointCloud cloud(3);
+
+  cloud.mutable_xyzs().transpose() << 0, 0, 0, 0, 0, 1, 0, 1, 1;
+
+  EXPECT_FALSE(cloud.has_normals());
+  cloud.EstimateNormals(10, 3);
+  EXPECT_TRUE(cloud.has_normals());
+
+  double kTol = 1e-6;
+  for (int i = 0; i < 3; ++i) {
+    CheckNormal(cloud.normal(i), Vector3f{1, 0, 0}, kTol);
+  }
+
+  cloud.mutable_xyzs().transpose() << 0, 0, 0, 1, 0, 0, 1, 0, 1;
+
+  cloud.EstimateNormals(10, 3);
+  for (int i = 0; i < 3; ++i) {
+    CheckNormal(cloud.normal(i), Vector3f{0, 1, 0}, kTol);
+  }
+
+  cloud.mutable_xyzs().transpose() << 0, 0, 0, 1, 0, 0, 0, 1, 1;
+  cloud.EstimateNormals(10, 3);
+  for (int i = 0; i < 3; ++i) {
+    CheckNormal(cloud.normal(i),
+                Vector3f{0, 1.0 / std::sqrt(2.0), -1.0 / std::sqrt(2.0)}, kTol);
+  }
+}
+
+// Test that points uniformly randomly distributed on the surface of the sphere
+// get normals close to the true normal of the sphere.
+GTEST_TEST(PointCloudTest, EstimateNormalsSphere) {
+  const int kSize{10000};
+  PointCloud cloud(kSize);
+
+  // Sample from a Gaussian and project it onto the sphere.
+  RandomGenerator generator(1234);
+  std::normal_distribution<double> distribution(0, 1.0);
+  for (int i = 0; i < 3 * kSize; ++i) {
+    cloud.mutable_xyzs().data()[i] = distribution(generator);
+  }
+  for (int i = 0; i < kSize; ++i) {
+    cloud.mutable_xyz(i).normalize();
+  }
+
+  cloud.EstimateNormals(0.1, 30);
+
+  double kTol = 1e-3;  // This will be loose unless kSize gets very large.
+  for (int i = 0; i < kSize; ++i) {
+    CheckNormal(cloud.normal(i), cloud.xyz(i), kTol);
+  }
+}
+
 }  // namespace
 }  // namespace perception
 }  // namespace drake

tools/workspace/default.bzl

@@ -63,6 +63,7 @@ load("@drake//tools/workspace/msgpack:repository.bzl", "msgpack_repository")
 load("@drake//tools/workspace/msgpack_lite_js:repository.bzl", "msgpack_lite_js_repository")  # noqa
 load("@drake//tools/workspace/mypy_extensions_internal:repository.bzl", "mypy_extensions_internal_repository")  # noqa
 load("@drake//tools/workspace/mypy_internal:repository.bzl", "mypy_internal_repository")  # noqa
+load("@drake//tools/workspace/nanoflann:repository.bzl", "nanoflann_repository")  # noqa
 load("@drake//tools/workspace/net_sf_jchart2d:repository.bzl", "net_sf_jchart2d_repository")  # noqa
 load("@drake//tools/workspace/nlopt:repository.bzl", "nlopt_repository")
 load("@drake//tools/workspace/nlopt_internal:repository.bzl", "nlopt_internal_repository")  # noqa
@@ -277,6 +278,8 @@ def add_default_repositories(excludes = [], mirrors = DEFAULT_MIRRORS):
         mypy_extensions_internal_repository(name = "mypy_extensions_internal", mirrors = mirrors)  # noqa
     if "mypy_internal" not in excludes:
         mypy_internal_repository(name = "mypy_internal", mirrors = mirrors)
+    if "nanoflann" not in excludes:
+        nanoflann_repository(name = "nanoflann", mirrors = mirrors)
     if "net_sf_jchart2d" not in excludes:
         net_sf_jchart2d_repository(name = "net_sf_jchart2d", mirrors = mirrors)
     if "nlopt" not in excludes:

tools/workspace/nanoflann/BUILD.bazel

@@ -0,0 +1,5 @@
+# -*- python -*-
+
+load("//tools/lint:lint.bzl", "add_lint_tests")
+
+add_lint_tests()

tools/workspace/nanoflann/package.BUILD.bazel

@@ -0,0 +1,24 @@
+# -*- python -*-
+
+load(
+    "@drake//tools/install:install.bzl",
+    "install",
+)
+
+licenses(["notice"])  # BSD
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "nanoflann",
+    hdrs = ["include/nanoflann.hpp"],
+    includes = ["include"],
+    deps = [],
+    linkstatic = 1,
+)
+
+# Install the license file.
+install(
+    name = "install",
+    docs = ["COPYING"],
+)

tools/workspace/nanoflann/repository.bzl

@@ -0,0 +1,15 @@
+# -*- python -*-
+
+load("@drake//tools/workspace:github.bzl", "github_archive")
+
+def nanoflann_repository(
+        name,
+        mirrors = None):
+    github_archive(
+        name = name,
+        repository = "jlblancoc/nanoflann",
+        commit = "v1.4.3",
+        sha256 = "cbcecf22bec528a8673a113ee9b0e134f91f1f96be57e913fa1f74e98e4449fa",  # noqa
+        build_file = ":package.BUILD.bazel",
+        mirrors = mirrors,
+    )