Add point cloud plane segmentation (#1287)

* Add point cloud plane segmentation

* Simplify the RANSAC implementation

* Add point cloud plane segmentation example

* Add function 'GetPlaneFromPoints'

* Add unit test

* Fix minor issues

* Fix Python formatting

Author: kongsgard

examples/Python/Basic/pointcloud_plane_segmentation.py

@@ -0,0 +1,28 @@
+# Open3D: www.open3d.org
+# The MIT License (MIT)
+# See license file or visit www.open3d.org for details
+
+# examples/Python/Basic/pointcloud_plane_segmentation.py
+
+import numpy as np
+import open3d as o3d
+
+if __name__ == "__main__":
+    pcd = o3d.io.read_point_cloud("../../TestData/fragment.pcd")
+
+    print(
+        "Find the plane model and the inliers of the largest planar segment in the point cloud."
+    )
+    plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
+                                             ransac_n=3,
+                                             num_iterations=250)
+
+    [a, b, c, d] = plane_model
+    print(f"Plane model: {a:.2f}x + {b:.2f}y + {c:.2f}z + {d:.2f} = 0")
+
+    inlier_cloud = pcd.select_down_sample(inliers)
+    inlier_cloud.paint_uniform_color([1.0, 0, 0])
+
+    outlier_cloud = pcd.select_down_sample(inliers, invert=True)
+
+    o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])

src/Open3D/Geometry/PointCloud.h

@@ -211,6 +211,20 @@ class PointCloud : public Geometry3D {
                                    size_t min_points,
                                    bool print_progress = false) const;
 
+    /// \brief Segment PointCloud plane using the RANSAC algorithm.
+    ///
+    /// \param distance_threshold Max distance a point can be from the plane
+    /// model, and still be considered an inlier.
+    /// \param ransac_n Number of initial points to be considered inliers in
+    /// each iteration.
+    /// \param num_iterations Number of iterations.
+    /// \return Returns the plane model ax + by + cz + d = 0 and the indices of
+    /// the plane inliers.
+    std::tuple<Eigen::Vector4d, std::vector<size_t>> SegmentPlane(
+            const double distance_threshold = 0.01,
+            const int ransac_n = 3,
+            const int num_iterations = 100) const;
+
     /// Factory function to create a pointcloud from a depth image and a camera
     /// model (PointCloudFactory.cpp)
     /// The input depth image can be either a float image, or a uint16_t image.

src/Open3D/Geometry/PointCloudSegmentation.cpp

@@ -0,0 +1,215 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// The MIT License (MIT)
+//
+// Copyright (c) 2019 www.open3d.org
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+// IN THE SOFTWARE.
+// ----------------------------------------------------------------------------
+
+#include "Open3D/Geometry/PointCloud.h"
+
+#include <Eigen/Dense>
+#include <algorithm>
+#include <iterator>
+#include <numeric>
+#include <random>
+#include <unordered_set>
+
+#include "Open3D/Geometry/TriangleMesh.h"
+#include "Open3D/Utility/Console.h"
+
+namespace open3d {
+namespace geometry {
+
+/// \class RANSACResult
+///
+/// \brief Stores the current best result in the RANSAC algorithm.
+class RANSACResult {
+public:
+    RANSACResult() : fitness_(0), inlier_rmse_(0) {}
+    ~RANSACResult() {}
+
+public:
+    double fitness_;
+    double inlier_rmse_;
+};
+
+// Calculates the number of inliers given a list of points and a plane model,
+// and the total distance between the inliers and the plane. These numbers are
+// then used to evaluate how well the plane model fits the given points.
+RANSACResult EvaluateRANSACBasedOnDistance(
+        const std::vector<Eigen::Vector3d> &points,
+        const Eigen::Vector4d plane_model,
+        std::vector<size_t> &inliers,
+        double distance_threshold,
+        double error) {
+    RANSACResult result;
+
+    for (size_t idx = 0; idx < points.size(); ++idx) {
+        Eigen::Vector4d point(points[idx](0), points[idx](1), points[idx](2),
+                              1);
+        double distance = std::abs(plane_model.dot(point));
+
+        if (distance < distance_threshold) {
+            error += distance;
+            inliers.emplace_back(idx);
+        }
+    }
+
+    size_t inlier_num = inliers.size();
+    if (inlier_num == 0) {
+        result.fitness_ = 0;
+        result.inlier_rmse_ = 0;
+    } else {
+        result.fitness_ = (double)inlier_num / (double)points.size();
+        result.inlier_rmse_ = error / std::sqrt((double)inlier_num);
+    }
+    return result;
+}
+
+// Find the plane such that the summed squared distance from the
+// plane to all points is minimized.
+//
+// Reference:
+// https://www.ilikebigbits.com/2015_03_04_plane_from_points.html
+Eigen::Vector4d GetPlaneFromPoints(const std::vector<Eigen::Vector3d> &points,
+                                   const std::vector<size_t> &inliers) {
+    Eigen::Vector3d centroid(0, 0, 0);
+    for (size_t idx : inliers) {
+        centroid += points[idx];
+    }
+    centroid /= double(inliers.size());
+
+    double xx = 0, xy = 0, xz = 0, yy = 0, yz = 0, zz = 0;
+
+    for (size_t idx : inliers) {
+        Eigen::Vector3d r = points[idx] - centroid;
+        xx += r(0) * r(0);
+        xy += r(0) * r(1);
+        xz += r(0) * r(2);
+        yy += r(1) * r(1);
+        yz += r(1) * r(2);
+        zz += r(2) * r(2);
+    }
+
+    double det_x = yy * zz - yz * yz;
+    double det_y = xx * zz - xz * xz;
+    double det_z = xx * yy - xy * xy;
+
+    Eigen::Vector3d abc;
+    if (det_x > det_y && det_x > det_z) {
+        abc = Eigen::Vector3d(det_x, xz * yz - xy * zz, xy * yz - xz * yy);
+    } else if (det_y > det_z) {
+        abc = Eigen::Vector3d(xz * yz - xy * zz, det_y, xy * xz - yz * xx);
+    } else {
+        abc = Eigen::Vector3d(xy * yz - xz * yy, xy * xz - yz * xx, det_z);
+    }
+
+    double norm = abc.norm();
+    // Return invalid plane if the points don't span a plane.
+    if (norm == 0) {
+        return Eigen::Vector4d(0, 0, 0, 0);
+    }
+    abc /= abc.norm();
+    double d = -abc.dot(centroid);
+    return Eigen::Vector4d(abc(0), abc(1), abc(2), d);
+}
+
+std::tuple<Eigen::Vector4d, std::vector<size_t>> PointCloud::SegmentPlane(
+        const double distance_threshold /* = 0.01 */,
+        const int ransac_n /* = 3 */,
+        const int num_iterations /* = 100 */) const {
+    RANSACResult result;
+    double error = 0;
+
+    // Initialize the plane model ax + by + cz + d = 0.
+    Eigen::Vector4d plane_model = Eigen::Vector4d(0, 0, 0, 0);
+    // Initialize the best plane model.
+    Eigen::Vector4d best_plane_model = Eigen::Vector4d(0, 0, 0, 0);
+
+    // Initialize consensus set.
+    std::vector<size_t> inliers;
+
+    size_t num_points = points_.size();
+    std::vector<size_t> indices(num_points);
+    std::iota(std::begin(indices), std::end(indices), 0);
+
+    std::random_device rd;
+    std::mt19937 rng(rd());
+
+    // Return if ransac_n is less than the required plane model parameters.
+    if (ransac_n < 3) {
+        utility::LogError(
+                "ransac_n should be set to higher than or equal to 3.");
+        return std::make_tuple(best_plane_model, inliers);
+    }
+
+    for (int itr = 0; itr < num_iterations; itr++) {
+        for (int i = 0; i < ransac_n; ++i) {
+            std::swap(indices[i], indices[rng() % num_points]);
+        }
+        inliers.clear();
+        for (int idx = 0; idx < ransac_n; ++idx) {
+            inliers.emplace_back(indices[idx]);
+        }
+
+        // Fit model to num_model_parameters randomly selected points among the
+        // inliers.
+        plane_model = TriangleMesh::ComputeTrianglePlane(
+                points_[inliers[0]], points_[inliers[1]], points_[inliers[2]]);
+        if (plane_model.isZero(0)) {
+            continue;
+        }
+
+        error = 0;
+        inliers.clear();
+        auto this_result = EvaluateRANSACBasedOnDistance(
+                points_, plane_model, inliers, distance_threshold, error);
+        if (this_result.fitness_ > result.fitness_ ||
+            (this_result.fitness_ == result.fitness_ &&
+             this_result.inlier_rmse_ < result.inlier_rmse_)) {
+            result = this_result;
+            best_plane_model = plane_model;
+        }
+    }
+
+    // Find the final inliers using best_plane_model.
+    inliers.clear();
+    for (size_t idx = 0; idx < points_.size(); ++idx) {
+        Eigen::Vector4d point(points_[idx](0), points_[idx](1), points_[idx](2),
+                              1);
+        double distance = std::abs(best_plane_model.dot(point));
+
+        if (distance < distance_threshold) {
+            inliers.emplace_back(idx);
+        }
+    }
+
+    // Improve best_plane_model using the final inliers.
+    best_plane_model = GetPlaneFromPoints(points_, inliers);
+
+    utility::LogDebug("RANSAC | Inliers: {:d}, Fitness: {:e}, RMSE: {:e}",
+                      inliers.size(), result.fitness_, result.inlier_rmse_);
+    return std::make_tuple(best_plane_model, inliers);
+}
+
+}  // namespace geometry
+}  // namespace open3d

src/Open3D/Registration/Registration.cpp

@@ -236,7 +236,7 @@ RegistrationResult RegistrationRANSACBasedOnCorrespondence(
             result = this_result;
         }
     }
-    utility::LogDebug("RANSAC: Fitness {:.4f}, RMSE {:.4f}", result.fitness_,
+    utility::LogDebug("RANSAC: Fitness {:e}, RMSE {:e}", result.fitness_,
                       result.inlier_rmse_);
     return result;
 }
@@ -368,7 +368,7 @@ RegistrationResult RegistrationRANSACBasedOnFeatureMatching(
     }
 #endif
     utility::LogDebug("total_validation : {:d}", total_validation);
-    utility::LogDebug("RANSAC: Fitness {:.4f}, RMSE {:.4f}", result.fitness_,
+    utility::LogDebug("RANSAC: Fitness {:e}, RMSE {:e}", result.fitness_,
                       result.inlier_rmse_);
     return result;
 }

src/Python/open3d_pybind/geometry/pointcloud.cpp

@@ -172,6 +172,10 @@ void pybind_pointcloud(py::module &m) {
                  "Spatial Databases with Noise', 1996. Returns a list of point "
                  "labels, -1 indicates noise according to the algorithm.",
                  "eps"_a, "min_points"_a, "print_progress"_a = false)
+            .def("segment_plane", &geometry::PointCloud::SegmentPlane,
+                 "Segments a plane in the point cloud using the RANSAC "
+                 "algorithm.",
+                 "distance_threshold"_a, "ransac_n"_a, "num_iterations"_a)
             .def_static(
                     "create_from_depth_image",
                     &geometry::PointCloud::CreateFromDepthImage,
@@ -294,6 +298,15 @@ void pybind_pointcloud(py::module &m) {
              {"min_points", "Minimum number of points to form a cluster."},
              {"print_progress",
               "If true the progress is visualized in the console."}});
+    docstring::ClassMethodDocInject(
+            m, "PointCloud", "segment_plane",
+            {{"distance_threshold",
+              "Max distance a point can be from the plane model, and still be "
+              "considered an inlier."},
+             {"ransac_n",
+              "Number of initial points to be considered inliers in each "
+              "iteration."},
+             {"num_iterations", "Number of iterations."}});
     docstring::ClassMethodDocInject(m, "PointCloud", "create_from_depth_image");
     docstring::ClassMethodDocInject(m, "PointCloud", "create_from_rgbd_image");
 }