-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata_prepare.py
154 lines (140 loc) · 6.36 KB
/
data_prepare.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
import os
import yaml
import pickle
import argparse
import numpy as np
from os.path import join, exists
from sklearn.neighbors import KDTree
from utils.data_process import DataProcessing as DP
def fnv_hash_vec(arr):
"""
FNV64-1A
"""
assert arr.ndim == 2
# Floor first for negative coordinates
arr = arr.copy()
arr = arr.astype(np.uint64, copy=False)
hashed_arr = np.uint64(14695981039346656037) * np.ones(arr.shape[0], dtype=np.uint64)
for j in range(arr.shape[1]):
hashed_arr *= np.uint64(1099511628211)
hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])
return hashed_arr
def ravel_hash_vec(arr):
"""
Ravel the coordinates after subtracting the min coordinates.
"""
assert arr.ndim == 2
arr = arr.copy()
arr -= arr.min(0)
arr = arr.astype(np.uint64, copy=False)
arr_max = arr.max(0).astype(np.uint64) + 1
keys = np.zeros(arr.shape[0], dtype=np.uint64)
# Fortran style indexing
for j in range(arr.shape[1] - 1):
keys += arr[:, j]
keys *= arr_max[j + 1]
keys += arr[:, -1]
return keys
def voxelize(coord, voxel_size=0.2, hash_type='fnv', mode=0):
discrete_coord = np.floor(coord / np.array(voxel_size))
if hash_type == 'ravel':
key = ravel_hash_vec(discrete_coord)
else:
key = fnv_hash_vec(discrete_coord)
idx_sort = np.argsort(key)
key_sort = key[idx_sort]
_, count = np.unique(key_sort, return_counts=True)
if mode == 0: # train mode
idx_select = np.cumsum(np.insert(count, 0, 0)[0:-1]) + np.random.randint(0, count.max(), count.size) % count
idx_unique = idx_sort[idx_select]
return idx_unique
else: # val mode
return idx_sort, count
def load_pc_kitti(pc_path):
scan = np.fromfile(pc_path, dtype=np.float32)
scan = scan.reshape((-1, 4))
points = scan[:, 0:3] # get xyz
intensity = scan[:, -1] # get intensity
return points, intensity
def load_label_kitti(label_path, remap_lut):
label = np.fromfile(label_path, dtype=np.uint32)
label = label.reshape((-1))
sem_label = label & 0xFFFF # semantic label in lower half
inst_label = label >> 16 # instance id in upper half
assert ((sem_label + (inst_label << 16) == label).all())
sem_label = remap_lut[sem_label]
return sem_label.astype(np.int32)
parser = argparse.ArgumentParser()
parser.add_argument('--src_path', default='/data/semantickitti', help='source dataset path [default: None]')
parser.add_argument('--dst_path', default='process_dataset', help='destination dataset path [default: None]')
parser.add_argument('--grid_size', type=float, default=0.06, help='Subsample Grid Size [default: 0.06]')
parser.add_argument('--yaml_config', default='utils/semantic-kitti.yaml', help='semantic-kitti.yaml path')
FLAGS = parser.parse_args()
data_config = FLAGS.yaml_config
DATA = yaml.safe_load(open(data_config, 'r'))
remap_dict = DATA["learning_map"]
max_key = max(remap_dict.keys())
remap_lut = np.zeros((max_key + 100), dtype=np.int32)
remap_lut[list(remap_dict.keys())] = list(remap_dict.values())
grid_size = FLAGS.grid_size
dataset_path = FLAGS.src_path
output_path = FLAGS.dst_path
seq_list = np.sort(os.listdir(dataset_path))
for seq_id in seq_list:
print('sequence' + seq_id + ' start')
seq_path = join(dataset_path, seq_id)
seq_path_out = join(output_path, seq_id)
pc_path = join(seq_path, 'velodyne')
pc_path_out = join(seq_path_out, 'velodyne')
KDTree_path_out = join(seq_path_out, 'KDTree')
intensity_path_out = join(seq_path_out, 'intensity')
os.makedirs(seq_path_out) if not exists(seq_path_out) else None
os.makedirs(pc_path_out) if not exists(pc_path_out) else None
os.makedirs(KDTree_path_out) if not exists(KDTree_path_out) else None
os.makedirs(intensity_path_out) if not exists(intensity_path_out) else None
if int(seq_id) < 11:
label_path = join(seq_path, 'labels')
label_path_out = join(seq_path_out, 'labels')
os.makedirs(label_path_out) if not exists(label_path_out) else None
scan_list = np.sort(os.listdir(pc_path))
for scan_id in scan_list:
print(scan_id)
points, intensity = load_pc_kitti(join(pc_path, scan_id))
labels = load_label_kitti(join(label_path, str(scan_id[:-4]) + '.label'), remap_lut)
index = voxelize(points, grid_size)
sub_points, sub_labels, sub_intensity = points[index], labels[index], intensity[index]
search_tree = KDTree(sub_points)
KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
np.save(join(pc_path_out, scan_id)[:-4], sub_points)
np.save(join(label_path_out, scan_id)[:-4], sub_labels)
np.save(join(intensity_path_out, scan_id)[:-4], sub_intensity)
with open(KDTree_save, 'wb') as f:
pickle.dump(search_tree, f)
if seq_id == '08':
proj_path = join(seq_path_out, 'proj')
os.makedirs(proj_path) if not exists(proj_path) else None
proj_inds = np.squeeze(search_tree.query(points, return_distance=False))
proj_inds = proj_inds.astype(np.int32)
proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl')
with open(proj_save, 'wb') as f:
pickle.dump([proj_inds], f)
else:
proj_path = join(seq_path_out, 'proj')
os.makedirs(proj_path) if not exists(proj_path) else None
scan_list = np.sort(os.listdir(pc_path))
for scan_id in scan_list:
print(scan_id)
points, intensity = load_pc_kitti(join(pc_path, scan_id))
index = voxelize(points, grid_size)
sub_points, sub_intensity = points[index], intensity[index]
search_tree = KDTree(sub_points)
proj_inds = np.squeeze(search_tree.query(points, return_distance=False))
proj_inds = proj_inds.astype(np.int32)
KDTree_save = join(KDTree_path_out, str(scan_id[:-4]) + '.pkl')
proj_save = join(proj_path, str(scan_id[:-4]) + '_proj.pkl')
np.save(join(pc_path_out, scan_id)[:-4], sub_points)
np.save(join(intensity_path_out, scan_id)[:-4], sub_intensity)
with open(KDTree_save, 'wb') as f:
pickle.dump(search_tree, f)
with open(proj_save, 'wb') as f:
pickle.dump([proj_inds], f)