Add axis calculation component for pose estimation

src/diffCheck.hh

@@ -6,6 +6,9 @@
 #include <loguru.hpp>
 
 #include <cilantro/cilantro.hpp>
+#include <cilantro/clustering/kmeans.hpp>
+
+#include <Eigen/Dense>
 
 // diffCheck includes
 #include "diffCheck/log.hh"

src/diffCheck/IOManager.cc

@@ -68,4 +68,11 @@ namespace diffCheck::io
         std::filesystem::path pathCloud = pathTestData / "test_pc_for_SOR_101pts_with_1_outlier.ply";
         return pathCloud.string();
     }
+
+    std::string GetTwoConnectedPlanesPlyPath()
+    {
+        std::filesystem::path pathTestData = GetTestDataDir();
+        std::filesystem::path pathCloud = pathTestData / "two_connected_planes_with_normals.ply";
+        return pathCloud.string();
+    }
 } // namespace diffCheck::io

src/diffCheck/IOManager.hh

@@ -42,4 +42,6 @@ namespace diffCheck::io
     std::string GetRoofQuarterPlyPath();
     /// @brief Get the path to the plane point cloud with one outlier
     std::string GetPlanePCWithOneOutliers();
+    /// @brief Get the path to the two connected planes ply test file
+    std::string GetTwoConnectedPlanesPlyPath();
 } // namespace diffCheck::io

src/diffCheck/geometry/DFPointCloud.cc

@@ -216,6 +216,51 @@ namespace diffCheck::geometry
             this->Normals.push_back(normal);
     }
 
+    std::vector<Eigen::Vector3d> DFPointCloud::GetPrincipalAxes(int nComponents)
+    {
+        std::vector<Eigen::Vector3d> principalAxes;
+
+        if (! this->HasNormals())
+        {
+            DIFFCHECK_WARN("The point cloud has no normals. Normals will be estimated with knn = 20.");
+            this->EstimateNormals(true, 20);
+        }
+
+        // Convert normals to Eigen matrix
+        Eigen::Matrix<double, 3, Eigen::Dynamic> normalMatrix(3, this->Normals.size());
+        for (size_t i = 0; i < this->Normals.size(); ++i)
+        {
+            normalMatrix.col(i) = this->Normals[i].cast<double>();
+        }
+
+        cilantro::KMeans<double, 3> kmeans(normalMatrix); 
+        kmeans.cluster(nComponents);
+
+        const cilantro::VectorSet3d& centroids = kmeans.getClusterCentroids();
+        const std::vector<size_t>& assignments = kmeans.getPointToClusterIndexMap();
+        std::vector<int> clusterSizes(nComponents, 0);
+        for (size_t i = 0; i < assignments.size(); ++i) 
+        {
+            clusterSizes[assignments[i]]++;
+        }
+        // Sort clusters by size
+        std::vector<std::pair<int, Eigen::Vector3d>> sortedClustersBySize(nComponents);
+        for (size_t i = 0; i < nComponents; ++i) 
+        {
+            sortedClustersBySize[i] = {clusterSizes[i], centroids.col(i)};
+        }
+        std::sort(sortedClustersBySize.begin(), sortedClustersBySize.end(), [](const auto& a, const auto& b) 
+        {
+            return a.first > b.first;
+        });
+
+        for(size_t i = 0; i < nComponents; ++i) 
+        {
+            principalAxes.push_back(sortedClustersBySize[i].second);
+        }
+        return principalAxes;
+    }
+
     void DFPointCloud::UniformDownsample(int everyKPoints)
     {
         auto O3DPointCloud = this->Cvt2O3DPointCloud();

src/diffCheck/geometry/DFPointCloud.hh

@@ -8,6 +8,8 @@
 
 #include <cilantro/utilities/point_cloud.hpp>
 #include <cilantro/core/nearest_neighbors.hpp>
+#include <cilantro/clustering/kmeans.hpp>
+
 
 namespace diffCheck::geometry
 {
@@ -89,6 +91,14 @@ namespace diffCheck::geometry
          */
         void RemoveStatisticalOutliers(int nbNeighbors, double stdRatio);
 
+        /**
+         * @brief Get the nCompoments principal axes of the normals of the point cloud
+         * It is used to compute the pose of "boxy" point clouds. It relies on KMeans clustering to find the main axes of the point cloud.
+         *  @param nComponents the number of components to compute (default 6, each of 3 main axes in both directions)
+         * @return std::vector<Eigen::Vector3d> the principal axes of the point cloud ordered by number of normals
+         */
+        std::vector<Eigen::Vector3d> GetPrincipalAxes(int nComponents = 6);
+
     public:  ///< Downsamplers
         /**
          * @brief Downsample the point cloud with voxel grid

src/diffCheckBindings.cc

@@ -55,6 +55,9 @@ PYBIND11_MODULE(diffcheck_bindings, m) {
         .def("remove_statistical_outliers", &diffCheck::geometry::DFPointCloud::RemoveStatisticalOutliers, 
             py::arg("nb_neighbors"), py::arg("std_ratio"))
 
+        .def("get_principal_axes", &diffCheck::geometry::DFPointCloud::GetPrincipalAxes,
+            py::arg("n_components") = 6)
+
         .def("load_from_PLY", &diffCheck::geometry::DFPointCloud::LoadFromPLY)
         .def("save_to_PLY", &diffCheck::geometry::DFPointCloud::SaveToPLY)
 

src/gh/components/DF_main_pc_axes/code.py

@@ -0,0 +1,69 @@
+#! python3
+
+from diffCheck import diffcheck_bindings
+from diffCheck import df_cvt_bindings
+from diffCheck import df_poses
+
+import Rhino
+
+from ghpythonlib.componentbase import executingcomponent as component
+
+import System
+
+def compute_dot_product(v1, v2):
+    """
+    Compute the dot product of two vectors.
+    """
+    return (v1.X * v2.X) + (v1.Y * v2.Y) + (v1.Z * v2.Z)
+
+class DFMainPCAxes(component):
+    def RunScript(self,
+            i_clouds: System.Collections.Generic.List[Rhino.Geometry.PointCloud],
+            i_reset: bool):
+
+        planes = []
+        all_poses_in_time = df_poses.DFPosesAssembly()
+        if i_reset:
+            all_poses_in_time.reset()
+            return None, None
+
+        previous_poses = all_poses_in_time.get_last_poses()
+        all_poses_this_time = []
+        for i, cloud in enumerate(i_clouds):
+            df_cloud = df_cvt_bindings.cvt_rhcloud_2_dfcloud(cloud)
+            if df_cloud is None:
+                return None, None
+            df_cloud.estimate_normals(True, 12)
+
+            df_points = df_cloud.get_axis_aligned_bounding_box()
+            df_point = (df_points[0] + df_points[1]) / 2
+            rh_point = Rhino.Geometry.Point3d(df_point[0], df_point[1], df_point[2])
+            vectors = []
+            # Get the main axes of the point cloud
+            previous_pose = previous_poses[i] if previous_poses else None
+            if previous_pose:
+                rh_previous_xDirection = Rhino.Geometry.Vector3d(previous_pose.xDirection[0], previous_pose.xDirection[1], previous_pose.xDirection[2])
+                rh_previous_yDirection = Rhino.Geometry.Vector3d(previous_pose.yDirection[0], previous_pose.yDirection[1], previous_pose.yDirection[2])
+                n_faces = all_poses_in_time.poses_per_element_dictionary[f"element_{i}"].n_faces
+            else:
+                rh_previous_xDirection = None
+                rh_previous_yDirection = None
+                n_faces = len(diffcheck_bindings.dfb_segmentation.DFSegmentation.segment_by_normal(df_cloud, 12, int(len(df_cloud.points)/20), True, int(len(df_cloud.points)/200), 1))
+
+            axes = df_cloud.get_principal_axes(n_faces)
+            for axe in axes:
+                vectors.append(Rhino.Geometry.Vector3d(axe[0], axe[1], axe[2]))
+
+            new_xDirection, new_yDirection = df_poses.select_vectors(vectors, rh_previous_xDirection, rh_previous_yDirection)
+
+            pose = df_poses.DFPose(
+                origin = [rh_point.X, rh_point.Y, rh_point.Z],
+                xDirection = [new_xDirection.X, new_xDirection.Y, new_xDirection.Z],
+                yDirection = [new_yDirection.X, new_yDirection.Y, new_yDirection.Z])
+            all_poses_this_time.append(pose)
+            plane = Rhino.Geometry.Plane(origin = rh_point, xDirection=new_xDirection, yDirection=new_yDirection)
+            planes.append(plane)
+
+        all_poses_in_time.add_step(all_poses_this_time)
+
+        return [planes, all_poses_in_time]

src/gh/components/DF_main_pc_axes/metadata.json

@@ -0,0 +1,60 @@
+{
+    "name": "DFPoseEstimation",
+    "nickname": "PoseEsimation",
+    "category": "diffCheck",
+    "subcategory": "PointCloud",
+    "description": "This compoment calculates the pose of a list of point clouds.",
+    "exposure": 4,
+    "instanceGuid": "22b0c6fc-bc16-4ff5-b789-e99776277f65",
+    "ghpython": {
+        "hideOutput": true,
+        "hideInput": true,
+        "isAdvancedMode": true,
+        "marshalOutGuids": true,
+        "iconDisplay": 2,
+        "inputParameters": [
+            {
+                "name": "i_clouds",
+                "nickname": "i_clouds",
+                "description": "clouds whose main axes are to be calculated",
+                "optional": true,
+                "allowTreeAccess": true,
+                "showTypeHints": true,
+                "scriptParamAccess": "list",
+                "wireDisplay": "default",
+                "sourceCount": 0,
+                "typeHintID": "pointcloud"
+            }, 
+            {
+                "name": "i_reset",
+                "nickname": "i_reset",
+                "description": "reset the history of the pose estimation",
+                "optional": true,
+                "allowTreeAccess": false,
+                "showTypeHints": true,
+                "scriptParamAccess": "item",
+                "wireDisplay": "default",
+                "sourceCount": 0,
+                "typeHintID": "bool"
+            }
+        ],
+        "outputParameters": [
+            {
+                "name": "o_planes",
+                "nickname": "o_planes",
+                "description": "The resulting planes of the pose estimation in the last iteration.",
+                "optional": false,
+                "sourceCount": 0,
+                "graft": false
+            },
+            {
+                "name": "o_history",
+                "nickname": "o_history",
+                "description": "The history of poses of all the elements.",
+                "optional": false,
+                "sourceCount": 0,
+                "graft": false
+            }
+        ]
+    }
+}   

src/gh/diffCheck/diffCheck/df_poses.py

@@ -0,0 +1,118 @@
+from scriptcontext import sticky as rh_sticky_dict
+import json
+from dataclasses import dataclass, field
+
+@dataclass
+class DFPose:
+    """
+    This class represents the pose of a single element at a given time in the assembly process.
+    """
+    origin: list
+    xDirection: list
+    yDirection: list
+
+@dataclass
+class DFPosesBeam:
+    """
+    This class contains the poses of a single beam, at different times in the assembly process.
+    It also contains the number of faces detected for this element, based on which the poses are calculated.
+    """
+    poses_dictionnary: dict
+    n_faces: int = 3
+
+    def add_pose(self, pose: DFPose, step_number: int):
+        """
+        Add a pose to the dictionary of poses.
+        """
+        self.poses_dictionnary[f"pose_{step_number}"] = pose
+
+    def set_n_faces(self, n_faces: int):
+        """
+        Set the number of faces detected for this element.
+        """
+        self.n_faces = n_faces
+
+@dataclass
+class DFPosesAssembly:
+    n_step: int = 0
+    poses_per_element_dictionary: dict = field(default_factory=lambda: rh_sticky_dict)
+
+    """
+    This class contains the poses of the different elements of the assembly, at different times in the assembly process.
+    """
+    def __post_init__(self):
+        """
+        Initialize the poses_per_element_dictionary with empty DFPosesBeam objects.
+        """
+        lengths = []
+        for element in self.poses_per_element_dictionary:
+            lengths.append(len(self.poses_per_element_dictionary[element].poses_dictionnary))
+        self.n_step = max(lengths) if lengths else 0
+
+    def add_step(self, new_poses: list[DFPose]):
+        for i, pose in enumerate(new_poses):
+            if f"element_{i}" not in self.poses_per_element_dictionary:
+                self.poses_per_element_dictionary[f"element_{i}"] = DFPosesBeam({}, 4)
+                for j in range(self.n_step):
+                    self.poses_per_element_dictionary[f"element_{i}"].add_pose(None, j)
+            self.poses_per_element_dictionary[f"element_{i}"].add_pose(pose, self.n_step)
+        self.n_step += 1
+
+    def get_last_poses(self):
+        """
+        Get the last poses of each element.
+        """
+        if self.n_step == 0:
+            return None
+        last_poses = []
+        for i in range(len(self.poses_per_element_dictionary)):
+            last_poses.append(self.poses_per_element_dictionary[f"element_{i}"].poses_dictionnary[f"pose_{self.n_step-1}"])
+        return last_poses
+
+    def reset(self):
+        """
+        Reset the assembly poses to the initial state.
+        """
+        self.n_step = 0
+        rh_sticky_dict.clear()
+
+    def save(self, file_path: str):
+        """
+        Save the assembly poses to a JSON file.
+        """
+        with open(file_path, 'w') as f:
+            json.dump(self.poses_per_element_dictionary, f, default=lambda o: o.__dict__, indent=4)
+
+
+def compute_dot_product(v1, v2):
+    """
+    Compute the dot product of two vectors.
+    """
+    return (v1.X * v2.X) + (v1.Y * v2.Y) + (v1.Z * v2.Z)
+
+
+def select_vectors(vectors, previous_xDirection, previous_yDirection):
+    """
+    Select the vectors that are aligned with the xDirection and yDirection.
+    """
+    if previous_xDirection is not None and previous_yDirection is not None:
+        sorted_vectors_by_alignment = sorted(vectors, key=lambda v: compute_dot_product(v, previous_xDirection), reverse=True)
+        new_xDirection = sorted_vectors_by_alignment[0]
+    else:
+        new_xDirection = vectors[0]
+
+    condidates_for_yDirection = []
+    for v in vectors:
+        if compute_dot_product(v, new_xDirection) ** 2 < 0.5:
+            condidates_for_yDirection.append(v)
+    if previous_xDirection is not None and previous_yDirection is not None:
+        sorted_vectors_by_perpendicularity = sorted(condidates_for_yDirection, key=lambda v: compute_dot_product(v, previous_yDirection), reverse=True)
+        new_xDirection = sorted_vectors_by_alignment[0]
+        new_yDirection = sorted_vectors_by_perpendicularity[0] - compute_dot_product(sorted_vectors_by_perpendicularity[0], new_xDirection) * new_xDirection
+        new_yDirection.Unitize()
+    else:
+        new_xDirection = vectors[0]
+        sorted_vectors = sorted(vectors[1:], key=lambda v: compute_dot_product(v, new_xDirection)**2)
+        new_yDirection = sorted_vectors[0] - compute_dot_product(vectors[1], new_xDirection) * new_xDirection
+        new_yDirection.Unitize()
+    return new_xDirection, new_yDirection

tests/unit_tests/DFPointCloudTest.cc

@@ -219,4 +219,12 @@ TEST_F(DFPointCloudTestFixture, Transform) {
 
 //-------------------------------------------------------------------------
 // Others
-//-------------------------------------------------------------------------
+//-------------------------------------------------------------------------
+
+TEST_F(DFPointCloudTestFixture, KMeansClusteringOfNormals) {
+    std::string path = diffCheck::io::GetTwoConnectedPlanesPlyPath();
+    diffCheck::geometry::DFPointCloud dfPointCloud2Planes;
+    dfPointCloud2Planes.LoadFromPLY(path);
+    std::vector<Eigen::Vector3d> axes = dfPointCloud2Planes.GetPrincipalAxes(2);
+    EXPECT_TRUE((axes[0] - Eigen::Vector3d(0, 0, 1)).norm() < 1e-2 || (axes[1] - Eigen::Vector3d(0, 0, 1)).norm() < 1e-2);
+}