-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathpreprocess.py
117 lines (90 loc) · 4 KB
/
preprocess.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
#!/usr/bin/env python
# -*- coding:UTF-8 -*-
# This code will be used to voxelize the input point cloud
import os
import multiprocessing
import numpy as np
data_dir = 'velodyne'
def process_pointcloud(point_cloud, cls='car'):
# Input:
# (N, 4)
# Output:
# voxel_dict
if cls == 'Car':
scene_size = np.array([4, 80, 70.4], dtype=np.float32)
voxel_size = np.array([0.4, 0.2, 0.2], dtype=np.float32)
grid_size = np.array([10, 400, 352], dtype=np.int64)
lidar_coord = np.array([0, 40, 3], dtype=np.float32)
max_point_number = 35
else:
scene_size = np.array([4, 40, 48], dtype=np.float32)
voxel_size = np.array([0.4, 0.2, 0.2], dtype=np.float32)
grid_size = np.array([10, 200, 240], dtype=np.int64)
lidar_coord = np.array([0, 20, 3], dtype=np.float32)
max_point_number = 45
np.random.shuffle(point_cloud)
shifted_coord = point_cloud[:, :3] + lidar_coord
# reverse the point cloud coordinate (X, Y, Z) -> (Z, Y, X)
voxel_index = np.floor(
shifted_coord[:, ::-1] / voxel_size).astype(np.int)
bound_x = np.logical_and(
voxel_index[:, 2] >= 0, voxel_index[:, 2] < grid_size[2])
bound_y = np.logical_and(
voxel_index[:, 1] >= 0, voxel_index[:, 1] < grid_size[1])
bound_z = np.logical_and(
voxel_index[:, 0] >= 0, voxel_index[:, 0] < grid_size[0])
bound_box = np.logical_and(np.logical_and(bound_x, bound_y), bound_z)
point_cloud = point_cloud[bound_box]
voxel_index = voxel_index[bound_box]
# [K, 3] coordinate buffer as described in the paper
coordinate_buffer = np.unique(voxel_index, axis=0)
K = len(coordinate_buffer)
T = max_point_number
# [K, 1] store number of points in each voxel grid
number_buffer = np.zeros(shape=(K), dtype=np.int64)
# [K, T, 7] feature buffer as described in the paper
feature_buffer = np.zeros(shape=(K, T, 7), dtype=np.float32)
# build a reverse index for coordinate buffer
index_buffer = {}
for i in range(K):
index_buffer[tuple(coordinate_buffer[i])] = i
for voxel, point in zip(voxel_index, point_cloud):
index = index_buffer[tuple(voxel)]
number = number_buffer[index]
if number < T:
feature_buffer[index, number, :4] = point
number_buffer[index] += 1
feature_buffer[:, :, -3:] = feature_buffer[:, :, :3] - \
feature_buffer[:, :, :3].sum(axis=1, keepdims=True)/number_buffer.reshape(K, 1, 1)
voxel_dict = {'feature_buffer': feature_buffer,
'coordinate_buffer': coordinate_buffer,
'number_buffer': number_buffer}
return voxel_dict
# def worker(filelist):
# for file in filelist:
# point_cloud = np.fromfile(
# os.path.join(data_dir, file), dtype=np.float32).reshape(-1, 4)
# name, extension = os.path.splitext(file)
# voxel_dict = process_pointcloud(point_cloud)
# output_dir = 'voxel' if cfg.DETECT_OBJ == 'Car' else 'voxel_ped'
# np.savez_compressed(os.path.join(output_dir, name), **voxel_dict)
def worker(filelist, cls='car'):
for file in filelist:
point_cloud = np.fromfile(
os.path.join(data_dir, file), dtype=np.float32).reshape(-1, 4)
name, extension = os.path.splitext(file)
voxel_dict = process_pointcloud(point_cloud)
output_dir = 'voxel' if cls == 'car' else 'voxel_ped'
np.savez_compressed(os.path.join(output_dir, name), **voxel_dict)
def preprocess(filelist):
for file in filelist:
point_cloud = np.fromfile(os.path.join(data_dir, file), dtype=np.float32).reshape(-1, 4)
# if __name__ == '__main__':
# filelist = [f for f in os.listdir(data_dir) if f.endswith('bin')]
# num_worker = 8
# for sublist in np.array_split(filelist, num_worker):
# p = multiprocessing.Process(target=worker, args=(sublist,))
# p.start()
if __name__=='__main__':
filelist = [f for f in os.listdir(data_dir) if f.endswith('bin')]
preprocess(filelist)