import numpy as np import open3d as o3d import torch import trimesh from scipy.spatial import cKDTree from utils.pose_util import PoseUtil class PtsUtil: @staticmethod def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005): o3d_pc = o3d.geometry.PointCloud() o3d_pc.points = o3d.utility.Vector3dVector(point_cloud) downsampled_pc = o3d_pc.voxel_down_sample(voxel_size) return np.asarray(downsampled_pc.points) @staticmethod def random_downsample_point_cloud(point_cloud, num_points, require_idx=False): if point_cloud.shape[0] == 0: if require_idx: return point_cloud, np.array([]) return point_cloud idx = np.random.choice(len(point_cloud), num_points, replace=True) if require_idx: return point_cloud[idx], idx return point_cloud[idx] @staticmethod def fps_downsample_point_cloud(point_cloud, num_points, require_idx=False): N = point_cloud.shape[0] mask = np.zeros(N, dtype=bool) sampled_indices = np.zeros(num_points, dtype=int) sampled_indices[0] = np.random.randint(0, N) distances = np.linalg.norm( point_cloud - point_cloud[sampled_indices[0]], axis=1 ) for i in range(1, num_points): farthest_index = np.argmax(distances) sampled_indices[i] = farthest_index mask[farthest_index] = True new_distances = np.linalg.norm( point_cloud - point_cloud[farthest_index], axis=1 ) distances = np.minimum(distances, new_distances) sampled_points = point_cloud[sampled_indices] if require_idx: return sampled_points, sampled_indices return sampled_points @staticmethod def random_downsample_point_cloud_tensor(point_cloud, num_points): idx = torch.randint(0, len(point_cloud), (num_points,)) return point_cloud[idx] @staticmethod def voxelize_points(points, voxel_size): voxel_indices = np.floor(points / voxel_size).astype(np.int32) unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True) return unique_voxels @staticmethod def transform_point_cloud(points, pose_mat): points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1) points_h = np.dot(pose_mat, points_h.T).T return points_h[:, :3] @staticmethod def get_overlapping_points( point_cloud_L, point_cloud_R, voxel_size=0.005, require_idx=False ): voxels_L, indices_L = PtsUtil.voxelize_points(point_cloud_L, voxel_size) voxels_R, _ = PtsUtil.voxelize_points(point_cloud_R, voxel_size) voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3) voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3) overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R) mask_L = np.isin( indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0] ) overlapping_points = point_cloud_L[mask_L] if require_idx: return overlapping_points, mask_L return overlapping_points @staticmethod def filter_points( points, points_normals, cam_pose, voxel_size=0.002, theta=45, z_range=(0.25, 0.5), ): """filter with z range""" points_cam = PtsUtil.transform_point_cloud(points, np.linalg.inv(cam_pose)) idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1]) z_filtered_points = points[idx] """ filter with normal """ sampled_points = PtsUtil.voxel_downsample_point_cloud( z_filtered_points, voxel_size ) kdtree = cKDTree(points_normals[:, :3]) _, indices = kdtree.query(sampled_points) nearest_points = points_normals[indices] normals = nearest_points[:, 3:] camera_axis = -cam_pose[:3, 2] normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True) cos_theta = np.dot(normals_normalized, camera_axis) theta_rad = np.deg2rad(theta) idx = cos_theta > np.cos(theta_rad) filtered_sampled_points = sampled_points[idx] return filtered_sampled_points[:, :3] @staticmethod 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) 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 = pipreg.registration_ransac_based_on_feature_matching( source_downsampled, model_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(4000000, 500), ) reg_icp = pipreg.registration_icp( source_downsampled, model_downsampled, voxel_size/2, reg_ransac.transformation, pipreg.TransformationEstimationPointToPlane(), 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 i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy") z = depth x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0] y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1] points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3) mask = mask.reshape(-1, 4) points_camera = np.concatenate( [points_camera, np.ones((points_camera.shape[0], 1))], axis=-1 ) points_world = np.dot(cam_extrinsic, points_camera.T).T[:, :3] data = { "points_world": points_world, "points_camera": points_camera, } return data