File size: 5,541 Bytes
f96995c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import numpy as np
import open3d as o3d
def filter_tabletop_points(pcd):
# RANSAC to find table plane
plane_model, inliers = pcd.segment_plane(distance_threshold=0.01, ransac_n=3, num_iterations=1000)
[a, b, c, d] = plane_model
table_plane = np.array([a, b, c, d])
# remove points on the table
pcd = pcd.select_by_index(inliers, invert=True)
outliers = None
new_outlier = None
rm_iter = 0
while new_outlier is None or len(new_outlier.points) > 0:
_, inlier_idx = pcd.remove_statistical_outlier(
nb_neighbors = 20, std_ratio = 1.5 + rm_iter * 0.5
)
new_pcd = pcd.select_by_index(inlier_idx)
new_outlier = pcd.select_by_index(inlier_idx, invert=True)
if outliers is None:
outliers = new_outlier
else:
outliers += new_outlier
pcd = new_pcd
rm_iter += 1
return pcd
def get_tabletop_points(rgb_list, depth_list, R_list, t_list, intr_list, bbox, depth_threshold=[0, 2]):
# increase if out of memory
stride = 1
pcd_all = o3d.geometry.PointCloud()
for i in range(len(rgb_list)):
intr = intr_list[i]
R_cam2board = R_list[i]
t_cam2board = t_list[i]
depth = depth_list[i].copy().astype(np.float32)
points = depth2fgpcd(depth, intr)
points = points.reshape(depth.shape[0], depth.shape[1], 3)
points = points[::stride, ::stride, :].reshape(-1, 3)
mask = np.logical_and(
(depth > depth_threshold[0]), (depth < depth_threshold[1])
) # (H, W)
mask = mask[::stride, ::stride].reshape(-1)
img = rgb_list[i].copy()
points = points[mask].reshape(-1, 3)
points = R_cam2board @ points.T + t_cam2board[:, None]
points = points.T # (N, 3)
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
colors = img[::stride, ::stride, :].reshape(-1, 3).astype(np.float64)
colors = colors[mask].reshape(-1, 3)
colors = colors[:, ::-1].copy()
pcd.colors = o3d.utility.Vector3dVector(colors / 255)
pcd_all += pcd
pcd = pcd_all
pcd = pcd.crop(o3d.geometry.AxisAlignedBoundingBox(min_bound=bbox[:, 0], max_bound=bbox[:, 1]))
pcd = pcd.voxel_down_sample(voxel_size=0.001)
# NOTE: simple filtering instead of instance-specific segmentation processing
pcd = filter_tabletop_points(pcd)
return pcd
def rpy_to_rotation_matrix(roll, pitch, yaw):
# Assume the input in in degree
roll = roll / 180 * np.pi
pitch = pitch / 180 * np.pi
yaw = yaw / 180 * np.pi
# Define the rotation matrices
Rx = np.array([[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]])
Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]])
Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]])
# Combine the rotations
R = Rz @ Ry @ Rx
return R
def rotation_matrix_to_rpy(rotation_matrix):
# Get the roll, pitch, yaw in radian
roll = np.arctan2(rotation_matrix[2, 1], rotation_matrix[2, 2])
pitch = np.arctan2(-rotation_matrix[2, 0], np.sqrt(rotation_matrix[2, 1] ** 2 + rotation_matrix[2, 2] ** 2))
yaw = np.arctan2(rotation_matrix[1, 0], rotation_matrix[0, 0])
# Get the roll, pitch, yaw in degree
roll = roll / np.pi * 180
pitch = pitch / np.pi * 180
yaw = yaw / np.pi * 180
return roll, pitch, yaw
def depth2fgpcd(depth, intrinsic_matrix):
H, W = depth.shape
x, y = np.meshgrid(np.arange(W), np.arange(H))
x = x.reshape(-1)
y = y.reshape(-1)
depth = depth.reshape(-1)
points = np.stack([x, y, np.ones_like(x)], axis=1)
points = points * depth[:, None]
points = points @ np.linalg.inv(intrinsic_matrix).T
return points
def similarity_transform(from_points, to_points):
assert len(from_points.shape) == 2, \
"from_points must be a m x n array"
assert from_points.shape == to_points.shape, \
"from_points and to_points must have the same shape"
N, m = from_points.shape
mean_from = from_points.mean(axis = 0)
mean_to = to_points.mean(axis = 0)
delta_from = from_points - mean_from # N x m
delta_to = to_points - mean_to # N x m
sigma_from = (delta_from * delta_from).sum(axis = 1).mean()
sigma_to = (delta_to * delta_to).sum(axis = 1).mean()
cov_matrix = delta_to.T.dot(delta_from) / N
U, d, V_t = np.linalg.svd(cov_matrix, full_matrices = True)
cov_rank = np.linalg.matrix_rank(cov_matrix)
S = np.eye(m)
if cov_rank >= m - 1: # and np.linalg.det(cov_matrix) < 0: # TODO touch calibration
S[m-1, m-1] = -1
elif cov_rank < m-1:
raise ValueError("colinearility detected in covariance matrix:\n{}".format(cov_matrix))
R = U.dot(S).dot(V_t)
c = (d * S.diagonal()).sum() / sigma_from
t = mean_to - c*R.dot(mean_from)
return c, R, t
def visualize_o3d(geometries):
frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.2, origin=[0, 0, 0])
geometries.append(frame)
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector()
viewer = o3d.visualization.Visualizer()
viewer.create_window()
for geometry in geometries:
viewer.add_geometry(geometry)
opt = viewer.get_render_option()
opt.background_color = np.asarray([1., 1., 1.])
viewer.run()
|