forked from jmwang0117/OccRWKV
-
Notifications
You must be signed in to change notification settings - Fork 0
/
semantic_kitti.py
217 lines (181 loc) · 8.97 KB
/
semantic_kitti.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
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
from glob import glob
import torch
import os
import yaml
import numpy as np
def mask_op(data, x_min, x_max):
mask = (data > x_min) & (data < x_max)
return mask
def get_mask(pc, lims):
mask_x = mask_op(pc[:, 0], lims[0][0] + 0.0001, lims[0][1] - 0.0001)
mask_y = mask_op(pc[:, 1], lims[1][0] + 0.0001, lims[1][1] - 0.0001)
mask_z = mask_op(pc[:, 2], lims[2][0] + 0.0001, lims[2][1] - 0.0001)
mask = (mask_x) & (mask_y) & mask_z
return mask
def unpack(compressed):
''' given a bit encoded voxel grid, make a normal voxel grid out of it. '''
uncompressed = np.zeros(compressed.shape[0] * 8, dtype=np.uint8)
uncompressed[::8] = compressed[:] >> 7 & 1
uncompressed[1::8] = compressed[:] >> 6 & 1
uncompressed[2::8] = compressed[:] >> 5 & 1
uncompressed[3::8] = compressed[:] >> 4 & 1
uncompressed[4::8] = compressed[:] >> 3 & 1
uncompressed[5::8] = compressed[:] >> 2 & 1
uncompressed[6::8] = compressed[:] >> 1 & 1
uncompressed[7::8] = compressed[:] & 1
return uncompressed
def augmentation_random_flip(data, flip_type, is_scan=False):
if flip_type==1:
if is_scan:
data[:, 0] = 51.2 - data[:, 0]
else:
data = np.flip(data, axis=0).copy()
elif flip_type==2:
if is_scan:
data[:, 1] = -data[:, 1]
else:
data = np.flip(data, axis=1).copy()
elif flip_type==3:
if is_scan:
data[:, 0] = 51.2 - data[:, 0]
data[:, 1] = -data[:, 1]
else:
data = np.flip(np.flip(data, axis=0), axis=1).copy()
return data
class SemanticKitti(torch.utils.data.Dataset):
CLASSES = ('unlabeled',
'car', 'bicycle', 'motorcycle', 'truck', 'other-vehicle',
'person', 'bicyclist', 'motorcyclist', 'road',
'parking', 'sidewalk', 'other-ground', 'building', 'fence',
'vegetation', 'trunk', 'terrain', 'pole', 'traffic-sign')
def __init__(self, data_root, data_config_file, setname,
lims,
sizes,
augmentation=False,
shuffle_index=False):
self.data_root = data_root
self.data_config = yaml.safe_load(open(data_config_file, 'r'))
self.sequences = self.data_config["split"][setname]
self.setname = setname
self.labels = self.data_config['labels']
self.learning_map = self.data_config["learning_map"]
self.learning_map_inv = self.data_config["learning_map_inv"]
self.color_map = self.data_config['color_map']
self.lims = lims
self.sizes = sizes
self.augmentation = augmentation
self.shuffle_index = shuffle_index
self.filepaths = {}
print(f"=> Parsing SemanticKITTI {self.setname}")
self.get_filepaths()
self.num_files_ = len(self.filepaths['occupancy'])
print("Using {} scans from sequences {}".format(self.num_files_, self.sequences))
print(f"Is aug: {self.augmentation}")
def get_filepaths(self,):
# fill in with names, checking that all sequences are complete
if self.setname != 'test':
for key in ['label_1_1', 'invalid_1_1', 'label_1_2', 'invalid_1_2', 'label_1_4', 'invalid_1_4', 'label_1_8', 'invalid_1_8', 'occluded', 'occupancy']:
self.filepaths[key] = []
else:
self.filepaths['occupancy'] = []
for seq in self.sequences:
# to string
seq = '{0:02d}'.format(int(seq))
print("parsing seq {}".format(seq))
if self.setname != 'test':
# Scale 1_1
self.filepaths['label_1_1'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.label')))
self.filepaths['invalid_1_1'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.invalid')))
# Scale 1_2
self.filepaths['label_1_2'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.label_1_2')))
self.filepaths['invalid_1_2'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.invalid_1_2')))
# Scale 1_4
self.filepaths['label_1_4'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.label_1_4')))
self.filepaths['invalid_1_4'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.invalid_1_4')))
# Scale 1_4
self.filepaths['label_1_8'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.label_1_8')))
self.filepaths['invalid_1_8'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.invalid_1_8')))
# occluded
self.filepaths['occluded'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.occluded')))
self.filepaths['occupancy'] += sorted(glob(os.path.join(self.data_root, seq, 'voxels', '*.bin')))
def get_data(self, idx, flip_type):
data_collection = {}
sc_remap_lut = self.get_remap_lut(completion=True)
ss_remap_lut = self.get_remap_lut()
for typ in self.filepaths.keys():
scale = int(typ.split('_')[-1]) if 'label' in typ or 'invalid' in typ else 1
if 'label' in typ:
scan_data = np.fromfile(self.filepaths[typ][idx], dtype=np.uint16)
if scale == 1:
scan_data = sc_remap_lut[scan_data]
else:
scan_data = unpack(np.fromfile(self.filepaths[typ][idx], dtype=np.uint8))
scan_data = scan_data.reshape((self.sizes[0]//scale, self.sizes[1]//scale, self.sizes[2]//scale))
scan_data = scan_data.astype(np.float32)
if self.augmentation:
scan_data = augmentation_random_flip(scan_data, flip_type)
data_collection[typ] = torch.from_numpy(scan_data)
points_path = self.filepaths['occupancy'][idx].replace('voxels', 'velodyne')
points = np.fromfile(points_path, dtype=np.float32)
points = points.reshape((-1, 4))
if self.setname != 'test':
points_label_path = self.filepaths['occupancy'][idx].replace('voxels', 'labels').replace('.bin', '.label')
points_label = np.fromfile(points_label_path, dtype=np.uint32)
points_label = points_label.reshape((-1))
points_label = points_label & 0xFFFF # semantic label in lower half
points_label = ss_remap_lut[points_label]
if self.shuffle_index:
pt_idx = np.random.permutation(np.arange(0, points.shape[0]))
points = points[pt_idx]
if self.setname != 'test':
points_label = points_label[pt_idx]
if self.lims:
filter_mask = get_mask(points, self.lims)
points = points[filter_mask]
if self.setname != 'test':
points_label = points_label[filter_mask]
if self.augmentation:
points = augmentation_random_flip(points, flip_type, is_scan=True)
data_collection['points'] = torch.from_numpy(points)
if self.setname != 'test':
data_collection['points_label'] = torch.from_numpy(points_label)
return data_collection
def __len__(self):
return self.num_files_
def get_n_classes(self):
return len(self.learning_map_inv)
def get_remap_lut(self, completion=False):
# put label from original values to xentropy
# or vice-versa, depending on dictionary values
# make learning map a lookup table
maxkey = max(self.learning_map.keys())
# +100 hack making lut bigger just in case there are unknown labels
remap_lut = np.zeros((maxkey + 100), dtype=np.int32)
remap_lut[list(self.learning_map.keys())] = list(self.learning_map.values())
# in completion we have to distinguish empty and invalid voxels.
# Important: For voxels 0 corresponds to "empty" and not "unlabeled".
if completion:
remap_lut[remap_lut == 0] = 255 # map 0 to 'invalid'
remap_lut[0] = 0 # only 'empty' stays 'empty'.
return remap_lut
def get_inv_remap_lut(self):
'''
remap_lut to remap classes of semantic kitti for training...
:return:
'''
# make lookup table for mapping
maxkey = max(self.learning_map_inv.keys())
# +100 hack making lut bigger just in case there are unknown labels
remap_lut = np.zeros((maxkey + 1), dtype=np.int32)
remap_lut[list(self.learning_map_inv.keys())] = list(self.learning_map_inv.values())
return remap_lut
def to_color(self, label):
# put label in original values
label = SemanticKitti.map(label, self.learning_map_inv)
# put label in color
return SemanticKitti.map(label, self.color_map)
def get_xentropy_class_string(self, idx):
return self.labels[self.learning_map_inv[idx]]
def __getitem__(self, idx):
flip_type = np.random.randint(0, 4) if self.augmentation else 0
return self.get_data(idx, flip_type), idx