308 lines
12 KiB
Python
308 lines
12 KiB
Python
|
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_old(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 voxel_downsample_point_cloud(point_cloud, voxel_size=0.005, require_idx=False):
|
||
|
|
voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
|
||
|
|
if require_idx:
|
||
|
|
_, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
|
||
|
|
idx_sort = np.argsort(inverse)
|
||
|
|
idx_unique = idx_sort[np.cumsum(counts)-counts]
|
||
|
|
downsampled_points = point_cloud[idx_unique]
|
||
|
|
return downsampled_points, idx_unique
|
||
|
|
else:
|
||
|
|
unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
|
||
|
|
return unique_voxels[0]*voxel_size
|
||
|
|
|
||
|
|
@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, normals, cam_pose, theta_limit=45, z_range=(0.2, 0.45)):
|
||
|
|
|
||
|
|
""" filter with normal """
|
||
|
|
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
|
||
|
|
cos_theta = np.dot(normals_normalized, np.array([0, 0, 1]))
|
||
|
|
theta = np.arccos(cos_theta) * 180 / np.pi
|
||
|
|
idx = theta < theta_limit
|
||
|
|
filtered_sampled_points = points[idx]
|
||
|
|
filtered_normals = normals[idx]
|
||
|
|
|
||
|
|
""" filter with z range """
|
||
|
|
points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose))
|
||
|
|
idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
|
||
|
|
z_filtered_points = filtered_sampled_points[idx]
|
||
|
|
z_filtered_normals = filtered_normals[idx]
|
||
|
|
return z_filtered_points[:, :3], z_filtered_normals
|
||
|
|
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def multi_scale_icp(
|
||
|
|
source, target, voxel_size_range, init_transformation=None, steps=20
|
||
|
|
):
|
||
|
|
pipreg = o3d.pipelines.registration
|
||
|
|
|
||
|
|
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 register_fine(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
|
||
|
|
)
|
||
|
|
)
|
||
|
|
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
|
||
|
|
)
|
||
|
|
)
|
||
|
|
reg_icp, icp_best_score = PtsUtil.multi_scale_icp(
|
||
|
|
source_downsampled,
|
||
|
|
model_downsampled,
|
||
|
|
voxel_size_range=(0.02, 0.001),
|
||
|
|
init_transformation=np.eye(4),
|
||
|
|
steps=50,
|
||
|
|
)
|
||
|
|
return reg_icp.transformation
|
||
|
|
@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
|