update

utils/pts_util.py

@@ -3,6 +3,7 @@ import open3d as o3d
 import torch
 import trimesh
 from scipy.spatial import cKDTree
+from utils.pose_util import PoseUtil
 
 
 class PtsUtil:
@@ -117,8 +118,8 @@ class PtsUtil:
         return filtered_sampled_points[:, :3]
 
     @staticmethod
-    def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.01):
-        radius_normal = voxel_size * 2
+    def old_register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.002):
+        radius_normal = voxel_size * 3
         pipreg = o3d.pipelines.registration
         model_pcd = o3d.geometry.PointCloud()
         model_pcd.points = o3d.utility.Vector3dVector(model.vertices)
@@ -152,7 +153,7 @@ class PtsUtil:
             source_fpfh,
             model_fpfh,
             mutual_filter=True,
-            max_correspondence_distance=voxel_size * 1.5,
+            max_correspondence_distance=voxel_size * 2,
             estimation_method=pipreg.TransformationEstimationPointToPoint(False),
             ransac_n=4,
             checkers=[pipreg.CorrespondenceCheckerBasedOnEdgeLength(0.9)],
@@ -162,14 +163,59 @@ class PtsUtil:
         reg_icp = pipreg.registration_icp(
             source_downsampled,
             model_downsampled,
-            voxel_size * 2,
+            voxel_size/2,
             reg_ransac.transformation,
             pipreg.TransformationEstimationPointToPlane(),
-            pipreg.ICPConvergenceCriteria(max_iteration=200),
+            pipreg.ICPConvergenceCriteria(max_iteration=2000),
         )
         
         return reg_icp.transformation
+    
+    @staticmethod
+    def chamfer_distance(pcl_a, pcl_b):
+        distances = np.linalg.norm(pcl_a[:, None] - pcl_b, axis=2)
+        min_distances = np.min(distances, axis=1)
+        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)
 
+        return best_pose
+    
     @staticmethod
     def get_pts_from_depth(depth, cam_intrinsic, cam_extrinsic):
         h, w = depth.shape