-
Notifications
You must be signed in to change notification settings - Fork 1
/
point_cloud_utils.py
179 lines (148 loc) · 5.94 KB
/
point_cloud_utils.py
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
168
169
170
171
172
173
174
175
176
177
178
179
# Transformation on point clouds
import numpy as np
import open3d as o3d
import copy
from tqdm import tqdm
import random
def generate_random_colors(N, seed=0):
colors = set() # Use a set to store unique colors
while len(colors) < N: # Keep generating colors until we have N unique ones
# Generate a random color and add it to the set
colors.add((random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
return list(colors) # Convert the set to a list before returning
def get_pcd(points: np.array):
"""
Convert numpy array to open3d point cloud
Args:
points: 3D points in camera coordinate [npoints, 3] or [npoints, 4]
Returns:
pcd: open3d point cloud
"""
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points[:, :3])
return pcd
def transform_pcd(points, T):
"""
Apply the perspective projection
Args:
points: 3D points in coordinate system 1 [npoints, 3]
T: Transformation matrix in homogeneous coordinates [4, 4]
Returns:
points: 3D points in coordinate system 2 [npoints, 3]
"""
pcd = get_pcd(points)
transformed_pcd = pcd.transform(T)
return np.asarray(transformed_pcd.points)
def filter_points_from_dict(points, filter_dict):
"""
Filter points based on a dict
Args:
points: 3D points in camera coordinate [npoints, 3]
filter_dict: dict that maps point indices to pixel coordinates
Returns:
points: 3D points in camera coordinate within image FOV [npoints, 3]
"""
inds = np.array(list(filter_dict.keys()))
return points[inds]
def filter_points_from_list(points, filter_list):
"""
Filter points based on a dict
Args:
points: 3D points in camera coordinate [npoints, 3]
filter_dict: dict that maps point indices to pixel coordinates
Returns:
points: 3D points in camera coordinate within image FOV [npoints, 3]
"""
inds = np.array(list(filter_list))
return points[inds]
def point_to_label(point_to_pixel: dict, label_map, label_is_color=True, label_is_instance=False):
'''
Args:
point_to_pixel: dict that maps point indices to pixel coordinates
label_map: label map of image
Returns:
point_to_label: dict that maps point indices to labels
'''
point_to_label_dict = {}
for index, point_data in point_to_pixel.items():
pixel = point_data['pixels']
if label_is_color:
color = label_map[pixel[1], pixel[0]]
if color.tolist() == [70,70,70]: # ignore unlabeled pixels
continue
point_to_label_dict[index] = tuple((color / 255))
elif label_is_instance:
instance_label = int(label_map[pixel[1], pixel[0]])
if instance_label: # ignore unlabeled pixels
point_to_label_dict[index] = instance_label
else:
continue
else: # label is continuous feature vector
point_to_label_dict[index] = label_map[pixel[1], pixel[0]]
return point_to_label_dict
def change_point_indices(point_to_X: dict, indices: list):
'''
Args:
point_to_X: dict that maps point indices to X
indices: list of indices
Returns:
point_to_X: dict that maps point indices to X
'''
point_to_X_new = {}
for old_index in point_to_X.keys():
new_index = indices[old_index]
point_to_X_new[new_index] = point_to_X[old_index]
return point_to_X_new
def transformation_matrix(rotation: np.array, translation: np.array):
"""
Create a transformation matrix from rotation and translation
Args:
rotation: Rotation matrix [3, 3]
translation: Translation vector [3, 1]
Returns:
T: Transformation matrix in homogeneous coordinates [4, 4]
"""
T = np.eye(4)
T[:3, :3] = rotation
T[:3, 3] = translation
return T
def kDTree_1NN_feature_reprojection(features_to, pcd_to, features_from, pcd_from, max_radius=None, no_feature_label=[1,0,0]):
'''
Args:
pcd_from: point cloud to be projected
pcd_to: point cloud to be projected to
search_method: search method ("radius", "knn")
search_param: search parameter (radius or k)
Returns:
features_to: features projected on pcd_to
'''
from_tree = o3d.geometry.KDTreeFlann(pcd_from)
for i, point in enumerate(np.asarray(pcd_to.points)):
[_, idx, _] = from_tree.search_knn_vector_3d(point, 1)
if max_radius is not None:
if np.linalg.norm(point - np.asarray(pcd_from.points)[idx[0]]) > max_radius:
features_to[i,:] = no_feature_label
else:
features_to[i,:] = features_from[idx[0]]
else:
features_to[i,:] = features_from[idx[0]]
return features_to
def intersect(pred_indices, gt_indices):
intersection = np.intersect1d(pred_indices, gt_indices)
return intersection.size / pred_indices.shape[0]
def remove_isolated_points(pcd, adjacency_matrix):
isolated_mask = ~np.all(adjacency_matrix == 0, axis=1)
adjacency_matrix = adjacency_matrix[isolated_mask][:, isolated_mask]
pcd = pcd.select_by_index(np.where(isolated_mask == True)[0])
return pcd, adjacency_matrix
def get_statistical_inlier_indices(pcd, nb_neighbors=20, std_ratio=2.0):
_, inlier_indices = copy.deepcopy(pcd).remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
return inlier_indices
def get_subpcd(pcd, indices, colors=False, normals=False):
subpcd = o3d.geometry.PointCloud()
subpcd.points = o3d.utility.Vector3dVector(np.asarray(pcd.points)[indices])
if colors:
subpcd.colors = o3d.utility.Vector3dVector(np.asarray(pcd.colors)[indices])
if normals:
subpcd.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals)[indices])
return subpcd