update

utils/pts_util.py

@@ -178,43 +178,143 @@ class PtsUtil:
         return np.sum(min_distances)
     
     @staticmethod
-    def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.008, max_iter=100000):
-        model_pts = model.vertices
-        sampled_world_pts = PtsUtil.voxel_downsample_point_cloud(pcl, voxel_size)
-        sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_size)
-        best_pose = np.eye(4)
-        best_pose[:3, 3] = np.mean(sampled_world_pts, axis=0) - np.mean(sampled_model_pts, axis=0)
-        best_distance = float('inf')
-        temperature = 1.0
-        cnt_unchange = 0
-        for _ in range(max_iter):
-            print(best_distance)
-            new_pose = best_pose.copy()
-            rotation_noise = np.random.randn(3)
-            rotation_noise /= np.linalg.norm(rotation_noise)
-            rotation_noise *= temperature
-            translation_noise = np.random.randn(3) * 0.1 * temperature
-            rotation_matrix = PoseUtil.get_uniform_rotation(0, 360)
-            new_pose[:3, :3] = rotation_matrix @ best_pose[:3, :3]
-            new_pose[:3, 3] += translation_noise
-            
-            distance = PtsUtil.chamfer_distance(
-                PtsUtil.transform_point_cloud(sampled_world_pts, new_pose),
-                sampled_model_pts
-            )
-            
-            if distance < best_distance:
-                best_pose, best_distance = new_pose, distance
-                cnt_unchange = 0
-            else:
-                cnt_unchange += 1
-                if cnt_unchange > 11110:
-                    break
-            
-            temperature *= 0.999
-            print(temperature)
+    def multi_scale_icp(
+        source, target, voxel_size_range, init_transformation=None, steps=20
+    ):
+        pipreg = o3d.pipelines.registration
 
-        return best_pose
+        if init_transformation is not None:
+            current_transformation = init_transformation
+        else:
+            current_transformation = np.identity(4)
+        cnt = 0
+        best_score = 1e10
+        best_reg = None
+        voxel_sizes = []
+        for i in range(steps):
+            voxel_sizes.append(
+                voxel_size_range[0]
+                + i * (voxel_size_range[1] - voxel_size_range[0]) / steps
+            )
+
+        for voxel_size in voxel_sizes:
+            radius_normal = voxel_size * 2
+            source_downsampled = source.voxel_down_sample(voxel_size)
+            source_downsampled.estimate_normals(
+                search_param=o3d.geometry.KDTreeSearchParamHybrid(
+                    radius=radius_normal, max_nn=30
+                )
+            )
+            target_downsampled = target.voxel_down_sample(voxel_size)
+            target_downsampled.estimate_normals(
+                search_param=o3d.geometry.KDTreeSearchParamHybrid(
+                    radius=radius_normal, max_nn=30
+                )
+            )
+
+            reg_icp = pipreg.registration_icp(
+                source_downsampled,
+                target_downsampled,
+                voxel_size * 2,
+                current_transformation,
+                pipreg.TransformationEstimationPointToPlane(),
+                pipreg.ICPConvergenceCriteria(max_iteration=500),
+            )
+            cnt += 1
+            if reg_icp.fitness == 0:
+                score = 1e10
+            else:
+                score = reg_icp.inlier_rmse / reg_icp.fitness
+            
+            if score < best_score:
+                best_score = score
+                best_reg = reg_icp
+
+        return best_reg, best_score
+
+    @staticmethod
+    def multi_scale_ransac(source_downsampled, target_downsampled, source_fpfh, model_fpfh, voxel_size_range, steps=20):
+        pipreg = o3d.pipelines.registration
+        cnt = 0
+        best_score = 1e10
+        best_reg = None
+        voxel_sizes = []
+        for i in range(steps):
+            voxel_sizes.append(
+                voxel_size_range[0]
+                + i * (voxel_size_range[1] - voxel_size_range[0]) / steps
+            )
+
+        for voxel_size in voxel_sizes:
+            reg_ransac = pipreg.registration_ransac_based_on_feature_matching(
+                source_downsampled,
+                target_downsampled,
+                source_fpfh,
+                model_fpfh,
+                mutual_filter=True,
+                max_correspondence_distance=voxel_size*2,
+                estimation_method=pipreg.TransformationEstimationPointToPoint(False),
+                ransac_n=4,
+                checkers=[pipreg.CorrespondenceCheckerBasedOnEdgeLength(0.9)],
+                criteria=pipreg.RANSACConvergenceCriteria(8000000, 500),
+            )
+            cnt += 1
+            if reg_ransac.fitness == 0:
+                score = 1e10
+            else:
+                score = reg_ransac.inlier_rmse / reg_ransac.fitness
+            if score < best_score:
+                best_score = score
+                best_reg = reg_ransac
+
+        return best_reg, best_score
+
+    @staticmethod
+    def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.01):
+        radius_normal = voxel_size * 2
+        pipreg = o3d.pipelines.registration
+        model_pcd = o3d.geometry.PointCloud()
+        model_pcd.points = o3d.utility.Vector3dVector(model.vertices)
+        model_downsampled = model_pcd.voxel_down_sample(voxel_size)
+        model_downsampled.estimate_normals(
+            search_param=o3d.geometry.KDTreeSearchParamHybrid(
+                radius=radius_normal, max_nn=30
+            )
+        )
+        model_fpfh = pipreg.compute_fpfh_feature(
+            model_downsampled,
+            o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
+        )
+
+        source_pcd = o3d.geometry.PointCloud()
+        source_pcd.points = o3d.utility.Vector3dVector(pcl)
+        source_downsampled = source_pcd.voxel_down_sample(voxel_size)
+        source_downsampled.estimate_normals(
+            search_param=o3d.geometry.KDTreeSearchParamHybrid(
+                radius=radius_normal, max_nn=30
+            )
+        )
+        source_fpfh = pipreg.compute_fpfh_feature(
+            source_downsampled,
+            o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
+        )
+
+        reg_ransac, ransac_best_score = PtsUtil.multi_scale_ransac(
+            source_downsampled,
+            model_downsampled,
+            source_fpfh,
+            model_fpfh,
+            voxel_size_range=(0.03, 0.005),
+            steps=3,
+        )
+        reg_icp, icp_best_score = PtsUtil.multi_scale_icp(
+            source_downsampled,
+            model_downsampled,
+            voxel_size_range=(0.02, 0.001),
+            init_transformation=reg_ransac.transformation,
+            steps=50,
+        )
+        return reg_icp.transformation
     
     @staticmethod
     def get_pts_from_depth(depth, cam_intrinsic, cam_extrinsic):