nbv_rec_control/register.py
2024-10-12 16:39:00 +08:00

110 lines
4.1 KiB
Python

import os
import numpy as np
import trimesh
from utils.pts_util import PtsUtil
import numpy as np
import open3d as o3d
import torch
import trimesh
from scipy.spatial import cKDTree
def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.002):
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 = 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_icp2 = pipreg.registration_icp(
source_downsampled,
model_downsampled,
voxel_size*10,
np.eye(4),
pipreg.TransformationEstimationPointToPlane(),
pipreg.ICPConvergenceCriteria(max_iteration=2000),
)
reg_icp = pipreg.registration_icp(
source_downsampled,
model_downsampled,
voxel_size,
reg_icp2.transformation,
pipreg.TransformationEstimationPointToPlane(),
pipreg.ICPConvergenceCriteria(max_iteration=2000),
)
return reg_icp2.transformation, reg_icp.transformation
if __name__ == "__main__":
model_dir = "/home/yan20/Desktop/nbv_rec/data/models/bear"
model_path = os.path.join(model_dir,"mesh.ply")
temp_name = "cad_model_world"
cad_model = trimesh.load(model_path)
pts_path = "/home/yan20/nbv_rec/project/franka_control/first_real_pts_bear.txt"
pts = np.loadtxt(pts_path)
very_coarse_real_to_cad = np.eye(4)
cad_model_pts = cad_model.vertices
very_coarse_real_to_cad[:3,3] = np.mean(cad_model_pts, axis=0) - np.mean(pts, axis=0)
very_coarse_cad_pts = PtsUtil.transform_point_cloud(pts, very_coarse_real_to_cad)
target_point = np.array([-3.422540776542676300e-02, -2.412379452948226755e-02, 1.123609126159126337e-01])
# 设置一个容忍度
tolerance = 1e-5
# 计算每个点与目标点之间的距离
distances = np.linalg.norm(very_coarse_cad_pts - target_point, axis=1)
# 统计距离小于容忍度的点的数量
count = np.sum(distances < tolerance)
print(count)
print(very_coarse_cad_pts.shape)
print(np.mean(pts, axis=0), np.mean(cad_model_pts, axis=0), np.mean(very_coarse_cad_pts, axis=0))
pts = pts[distances > tolerance]
np.savetxt(os.path.join(temp_name + "_filtered.txt"), pts)
# very_coarse_real_to_cad[:3,3] = np.mean(cad_model_pts, axis=0) - np.mean(pts, axis=0)
# very_coarse_cad_pts = PtsUtil.transform_point_cloud(pts, very_coarse_real_to_cad)
# np.savetxt(os.path.join(temp_name + "_very_coarse_reg.txt"), very_coarse_cad_pts)
# real_to_cad = PtsUtil.register(very_coarse_cad_pts, cad_model)
# cad_pts = PtsUtil.transform_point_cloud(very_coarse_cad_pts, real_to_cad)
# np.savetxt(os.path.join(temp_name + "_reg.txt"), cad_pts)