-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathmerge_lidar.py
127 lines (105 loc) · 4.6 KB
/
merge_lidar.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
import numpy as np
import os, shutils, json
from plyfile import PlyData, PlyElement
from scipy.spatial.transform import Rotation as R
import math
def pose_unreal2opengl(c2w_mat):
translation = c2w_mat[:3, 3]
rot = R.from_matrix(c2w_mat[:3, :3])
rot_vec = rot.as_rotvec()
rot_vec_new = rot_vec[[1, 2, 0]]
rot_vec_new[0] *= -1
rot_vec_new[2] *= -1
rot = R.from_rotvec(rot_vec_new)
translation_new = translation[[1, 2, 0]]
translation_new[1] *= -1
c2w_mat = np.eye(4)
c2w_mat[:3, :3] = rot.as_matrix()
c2w_mat[:3, 3] = translation_new
rot = np.eye(4)
# rot[1,1]=-1
# rot[2, 2] = -1
c2w_mat = c2w_mat @ rot
return c2w_mat
# 读取PLY文件
def read_ply(file_path):
with open(file_path, 'rb') as f:
plydata = PlyData.read(f)
vertex_data = plydata['vertex']
points = np.vstack([vertex_data['x'], vertex_data['y'], vertex_data['z']]).T
return points
# 将点云应用于位姿
def transform_point_cloud(points, pose):
transformed_points = np.hstack([points, np.ones((points.shape[0], 1))])
transformed_points = np.dot(transformed_points, pose[:3, :].T)
return transformed_points[:, :3]
# 合并点云
def merge_point_clouds(point_clouds):
merged_cloud = np.vstack(point_clouds)
return merged_cloud
# 保存PLY文件
def save_ply(xyz, file_path):
if 'txt' in file_path:
rgb = np.random.randint(0, 255, size=(xyz.shape[0], 3))
with open(file_path, 'w') as f:
for i in range(len(xyz)):
f.write(f"{int(i)} {xyz[i][0]} {xyz[i][1]} {xyz[i][2]} {rgb[i][0]} {rgb[i][1]} {rgb[i][2]}\n")
else:
# Define the dtype for the structured array
dtype = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'),
('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]
normals = np.zeros_like(xyz)
rgb = np.random.random((xyz.shape[0], 3))
elements = np.empty(xyz.shape[0], dtype=dtype)
attributes = np.concatenate((xyz, normals, rgb), axis=1)
elements[:] = list(map(tuple, attributes))
# Create the PlyData object and write to file
el = PlyElement.describe(elements, 'vertex')
PlyData([el]).write(file_path)
def generate_point_cloud_text(scene_path):
start_timestep, end_timestep = 5, 155
point_cloud_files = []
poses = []
for t in range(start_timestep, end_timestep):
# for t in range(start_timestep, end_timestep, 5):
with open(os.path.join(scene_path, "train_pic", f"train_camera_extrinsics_{t:06d}.json"), 'r') as file:
ego_to_world_current = np.array(json.load(file)['transform_matrix'])
lidar_to_world = pose_unreal2opengl(ego_to_world_current)
# lidar_to_world = ego_to_world_current
lidar_file = os.path.join(scene_path, "train_pic", f"train_lidar_{t:05d}.ply")
poses.append(lidar_to_world)
point_cloud_files.append(lidar_file)
# transfer point cloud
point_clouds = []
for file, pose in zip(point_cloud_files, poses):
points = read_ply(file)
opengl_point_cloud = points.copy()
opengl_point_cloud[:, 0] = points[:, 2]
opengl_point_cloud[:, 1] = points[:, 1]
opengl_point_cloud[:, 2] = points[:, 0]
radian_z = -90 * np.pi / 180 # rotate around Z axis
Rotation_Matrix_1 = np.array([
[math.cos(radian_z), -math.sin(radian_z), 0],
[math.sin(radian_z), math.cos(radian_z), 0],
[0, 0, 1]])
opengl_point_cloud = np.dot(Rotation_Matrix_1, opengl_point_cloud.T).T
transformed_points = transform_point_cloud(opengl_point_cloud, pose)
point_clouds.append(transformed_points)
# merge point cloud
merged_cloud = merge_point_clouds(point_clouds)
# Save the merged point cloud
save_ply(merged_cloud, f'{scene_path.split("/")[-1]+"_cam1"}/merged_cloud.ply')
save_ply(merged_cloud, f'{scene_path.split("/")[-1]+"_cam3"}/merged_cloud.ply')
# generate_point_cloud_text('data/carla_pic_0603_Town01')
# print("Done processing carla_pic_0603_Town01.")
# generate_point_cloud_text('data/carla_pic_0603_Town02')
# print("Done processing carla_pic_0603_Town02.")
# generate_point_cloud_text('data/carla_pic_0603_Town03')
# print("Done processing carla_pic_0603_Town03.")
# generate_point_cloud_text('data/carla_pic_0603_Town04')
# print("Done processing carla_pic_0603_Town04.")
# generate_point_cloud_text('data/carla_pic_0603_Town05')
# print("Done processing carla_pic_0603_Town05.")
generate_point_cloud_text('data/carla_pic_0603_Town10')
print("Done processing carla_pic_0603_Town10.")