def load_mode(self, mode):
indices = self.mode_indices[mode]
num_to_load = len(indices)
ret = []
for load_index in range(num_to_load):
labels, tag, voxel, rgb, raw_lidar = [], [], [], [], []
vox_feature, vox_number, vox_coordinate = [], [], []
rgb.append(
cv2.resize(cv2.imread(self.f_rgb[load_index]),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT)))
raw_lidar.append(
np.fromfile(self.f_lidar[load_index],
dtype=np.float32).reshape((-1, 4)))
labels.append([
line
for line in open(self.f_label[load_index], 'r').readlines()
])
tag.append(self.data_tag[load_index])
voxel.append(process_pointcloud(raw_lidar[-1]))
bsize, per_vox_feat, per_vox_num, per_vox_coor = build_input(
[voxel[-1]])
vox_feature.append(per_vox_feat)
vox_number.append(per_vox_num)
vox_coordinate.append(per_vox_coor)
ret.append((np.array(tag), np.array(labels), np.array(vox_feature),
np.array(vox_number), np.array(vox_coordinate),
np.array(rgb), np.array(raw_lidar)))
return ret
def __call__(self, load_index):
if self.aug:
ret = aug_data(self.data_tag[load_index],
self.data_dir,
sampler=self.db_sampler)
else:
rgb = cv2.resize(cv2.imread(self.f_rgb[load_index]),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT))
raw_lidar = np.fromfile(self.f_lidar[load_index],
dtype=np.float32).reshape((-1, 4))
if not self.is_testset:
labels = [
line for line in open(self.f_label[load_index],
'r').readlines()
]
else:
labels = ['']
tag = self.data_tag[load_index]
if self.f_voxel is None:
voxel = process_pointcloud(tag, raw_lidar)
else:
voxel_files = np.load(self.f_voxel[load_index])
voxel = {}
voxel['feature_buffer'] = voxel_files['feature_buffer']
voxel['coordinate_buffer'] = voxel_files['coordinate_buffer']
voxel['number_buffer'] = voxel_files['number_buffer']
voxel['mask_buffer'] = voxel_files['mask_buffer']
ret = [tag, rgb, raw_lidar, voxel, labels]
return ret
def __call__(self, load_index):
if aug:
ret = aug_data(self.data_tag[load_index], self.data_dir)
tag.append(ret[0])
rgb.append(ret[1])
raw_lidar.append(ret[2])
voxel.append(ret[3])
labels.append(ret[4])
else:
rgb.append(
cv2.resize(cv2.imread(f_rgb[load_index]),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT)))
raw_lidar.append(
np.fromfile(f_lidar[load_index], dtype=np.float32).reshape(
(-1, 4)))
if not is_testset:
labels.append([
line
for line in open(f_label[load_index], 'r').readlines()
])
else:
labels.append([''])
tag.append(data_tag[load_index])
voxel.append(process_pointcloud(raw_lidar[-1]))
return n
def __call__(self, load_index):
if self.aug:
ret = aug_data(self.data_tag[load_index], self.data_dir)
else:
#rgb = cv2.resize(cv2.imread(self.f_rgb[load_index]), (cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT))
rgb = 0 # placeholder
rgb = cv2.resize(
cv2.imread(
'/home/ziyanw1/data/KITTI/training/image_2/001201.png'),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT))
#rgb.append( cv2.imread(f_rgb[load_index]) )
raw_lidar = np.fromfile(self.f_lidar[load_index],
dtype=np.float32).reshape((-1, 4))
raw_obj_pc = np.fromfile(self.f_obj[load_index],
dtype=np.float32).reshape((-1, 4))
if not self.is_testset:
labels = [
line for line in open(self.f_label[load_index],
'r').readlines()
]
else:
labels = ['']
tag = self.data_tag[load_index]
voxel = process_pointcloud(raw_lidar)
ret = [tag, rgb, raw_lidar, raw_obj_pc, voxel, labels]
return ret
def run(self):
raw_lidar = self.np_p_ranged
# print(raw_lidar.shape) # DEBUG
voxel = process_pointcloud(raw_lidar)
# print("len of voxel['feature_buffer]: {}".format(len(voxel['feature_buffer']))) # 3000~5000
# print("len of voxel['coordinate_buffer']: {}".format(len(voxel['coordinate_buffer']))) # same as above
# print("len of voxel['number_buffer']: {}".format(len(voxel['number_buffer']))) #
return raw_lidar, voxel
def fill_queue(self, batch_size=0):
load_index = self.load_index
self.load_index += batch_size
if self.load_index >= self.dataset_size:
if not self.is_testset: # test set just end
if self.require_shuffle:
self.shuffle_dataset()
load_index = 0
self.load_index = load_index + batch_size
else:
self.work_exit.value = True
labels, tag, voxel, rgb, raw_lidar = [], [], [], [], []
for _ in range(batch_size):
try:
if self.aug:
ret = aug_data(self.data_tag[load_index], self.object_dir)
tag.append(ret[0])
rgb.append(ret[1])
raw_lidar.append(ret[2])
voxel.append(ret[3])
labels.append(ret[4])
else:
rgb.append(cv2.resize(cv2.imread(
self.f_rgb[load_index]), (cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT)))
raw_lidar.append(np.fromfile(
self.f_lidar[load_index], dtype=np.float32).reshape((-1, 4)))
if not self.is_testset:
labels.append([line for line in open(
self.f_label[load_index], 'r').readlines()])
else:
#labels.append(['Dummy 0.00 0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0'])
labels.append([''])
tag.append(self.data_tag[load_index])
voxel.append(process_pointcloud(raw_lidar[-1]))
load_index += 1
except:
if not self.is_testset: # test set just end
self.load_index = 0
if self.require_shuffle:
self.shuffle_dataset()
else:
self.work_exit.value = True
# only for voxel -> [gpu, k_single_batch, ...]
vox_feature, vox_number, vox_coordinate = [], [], []
single_batch_size = int(self.batch_size / self.multi_gpu_sum)
for idx in range(self.multi_gpu_sum):
_, per_vox_feature, per_vox_number, per_vox_coordinate = build_input(
voxel[idx * single_batch_size:(idx + 1) * single_batch_size])
vox_feature.append(per_vox_feature)
vox_number.append(per_vox_number)
vox_coordinate.append(per_vox_coordinate)
self.dataset_queue.put_nowait(
(labels, (vox_feature, vox_number, vox_coordinate), rgb, raw_lidar, tag))
def gen_batch(f_lidar):
lidar = np.fromfile(f_lidar, dtype=np.float32).reshape((-1, 4))
voxel_dict = process_pointcloud('pc', lidar)
num_voxels = voxel_dict['feature_buffer'].shape[0]
if num_voxels < batch_size:
return None, None
index = np.random.choice(num_voxels, batch_size, replace=False)
batch_pc = voxel_dict['feature_buffer'][index, :, :3]
batch_mask = voxel_dict['mask_buffer'][index]
return batch_pc, batch_mask
def _load_data(self, index, is_testset = False):
file_id = self.indices[index]
rgb = cv2.imread(self.f_rgb[file_id])
raw_lidar = np.fromfile(self.f_lidar[file_id], dtype = np.float32).reshape((-1, 4))
if not is_testset:
labels = [line for line in open(self.f_label[file_id], 'r').readlines()]
else:
labels = ['']
tag = self.data_tag[index]
voxel = process_pointcloud(raw_lidar)
return tag, rgb, raw_lidar, voxel, labels
def sample_test_data(data_dir, batch_size=1, multi_gpu_sum=1):
f_rgb = glob.glob(os.path.join(data_dir, 'image_2', '*.png'))
f_lidar = glob.glob(os.path.join(data_dir, 'velodyne', '*.bin'))
f_label = glob.glob(os.path.join(data_dir, 'label_2', '*.txt'))
f_rgb.sort()
f_lidar.sort()
f_label.sort()
data_tag = [name.split('/')[-1].split('.')[-2] for name in f_rgb]
assert (len(data_tag) == len(f_rgb) ==
len(f_lidar)), "dataset folder is not correct"
nums = len(f_rgb)
indices = list(range(nums))
np.random.shuffle(indices)
num_batches = int(math.floor(nums / float(batch_size)))
excerpt = indices[0:batch_size]
labels, tag, voxel, rgb, raw_lidar = [], [], [], [], []
for load_index in excerpt:
rgb.append(
cv2.resize(cv2.imread(f_rgb[load_index]),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT)))
#rgb.append( cv2.imread(f_rgb[load_index]) )
raw_lidar.append(
np.fromfile(f_lidar[load_index], dtype=np.float32).reshape(
(-1, 4)))
labels.append(
[line for line in open(f_label[load_index], 'r').readlines()])
tag.append(data_tag[load_index])
voxel.append(process_pointcloud(raw_lidar[-1]))
# only for voxel -> [gpu, k_single_batch, ...]
vox_feature, vox_number, vox_coordinate = [], [], []
single_batch_size = int(batch_size / multi_gpu_sum)
for idx in range(multi_gpu_sum):
_, per_vox_feature, per_vox_number, per_vox_coordinate = build_input(
voxel[idx * single_batch_size:(idx + 1) * single_batch_size])
vox_feature.append(per_vox_feature)
vox_number.append(per_vox_number)
vox_coordinate.append(per_vox_coordinate)
ret = (np.array(tag), np.array(labels), np.array(vox_feature),
np.array(vox_number), np.array(vox_coordinate), np.array(rgb),
np.array(raw_lidar))
return ret
def __call__(self,load_index):
if self.aug:
ret = aug_data(self.data_tag[load_index], self.data_dir)
else:
rgb = cv2.resize(cv2.imread(self.f_rgb[load_index]), (cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT))
raw_lidar = np.fromfile(self.f_lidar[load_index], dtype=np.float32).reshape((-1, 4))
if not self.is_testset:
labels = [line for line in open(self.f_label[load_index], 'r').readlines()]
else:
labels = ['']
tag = self.data_tag[load_index]
voxel = process_pointcloud(raw_lidar)
ret = [tag, rgb, raw_lidar, voxel, labels]
return ret
def load(self, batch_size=1):
if self.load_index >= self.dataset_size:
return None
num_to_load = batch_size
labels, tag, voxel, rgb, raw_lidar = [], [], [], [], []
vox_feature, vox_number, vox_coordinate = [], [], []
# Load normal data
for i in range(batch_size):
load_index = self.load_index + i
rgb.append(
cv2.resize(cv2.imread(self.f_rgb[load_index]),
(cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT)))
raw_lidar.append(
np.fromfile(self.f_lidar[load_index],
dtype=np.float32).reshape((-1, 4)))
if not self.is_testset:
labels.append([
line for line in open(self.f_label[load_index],
'r').readlines()
])
else:
labels.append([''])
tag.append(self.data_tag[load_index])
voxel.append(process_pointcloud(raw_lidar[-1]))
bsize, per_vox_feat, per_vox_num, per_vox_coor = build_input(
[voxel[-1]])
vox_feature.append(per_vox_feat)
vox_number.append(per_vox_num)
vox_coordinate.append(per_vox_coor)
#print('bsize', bsize)
#print('vox_feature', np.array(vox_feature).shape)
#print('vox_number', vox_number)
# Create augmented data
augIdxs, aug_voxel = [], []
if self.aug and not self.is_testset:
num_to_load += self.aug_num
# Get the indices of members to be augmented
augIdxs = np.random.choice(np.arange(
self.load_index, self.load_index + batch_size + 1),
self.aug_num,
replace=False)
for idx in augIdxs:
ret = aug_data(self.data_tag[idx], self.object_dir)
tag.append(ret[0])
rgb.append(ret[1])
raw_lidar.append(ret[2])
aug_voxel.append(ret[3])
labels.append(ret[4])
bsize, per_vox_feat, per_vox_coor, per_vox_num = build_input(
[aug_voxel[-1]])
vox_feature.append(per_vox_feat)
#print('per_vox_feat', np.array(per_vox_feat).shape)
vox_number.append(per_vox_num)
vox_coordinate.append(per_vox_coor)
#print('bsize', bsize)
#print('vox_feature', np.array(vox_feature).shape)
ret = (np.array(tag), np.array(labels), np.array(vox_feature),
np.array(vox_number), np.array(vox_coordinate), np.array(rgb),
np.array(raw_lidar))
flag = False
self.load_index += batch_size
if self.load_index >= self.dataset_size:
flag = True
return flag, ret
def aug_data(tag, object_dir, aug_pc=True, use_newtag=False, sampler=None):
np.random.seed()
rgb = cv2.imread(os.path.join(object_dir, 'image_2', tag + '.png'))
assert rgb is not None, print('ERROR rgb {} {}'.format(object_dir, tag))
rgb = cv2.resize(rgb, (cfg.IMAGE_WIDTH, cfg.IMAGE_HEIGHT))
#
lidar_path = os.path.join(object_dir, cfg.VELODYNE_DIR, tag + '.bin')
lidar = np.fromfile(lidar_path, dtype=np.float32).reshape(-1, 4)
assert lidar.shape[0], print('ERROR lidar {} {}'.format(object_dir, tag))
#
label_path = os.path.join(object_dir, 'label_2', tag + '.txt')
label = np.array([line
for line in open(label_path, 'r').readlines()]) # (N')
classes = np.array([line.split()[0] for line in label]) # (N')
# (N', 7) x, y, z, h, w, l, r
gt_box3d = label_to_gt_box3d([tag],
np.array(label)[np.newaxis, :],
cls=cfg.DETECT_OBJ,
coordinate='lidar')
gt_box3d = gt_box3d[0]
if sampler is not None:
avoid_collision_boxes = gt_box3d.copy()
avoid_collision_boxes[:, 3] = gt_box3d[:, 5]
avoid_collision_boxes[:, 4] = gt_box3d[:, 4]
avoid_collision_boxes[:, 5] = gt_box3d[:, 3]
sampled_dict = sampler.sample_all(cfg.DETECT_OBJ,
avoid_collision_boxes)
if sampled_dict is not None:
sampled_box = sampled_dict["gt_boxes"].copy()
sampled_box[:, 3] = sampled_dict["gt_boxes"][:, 5]
sampled_box[:, 4] = sampled_dict["gt_boxes"][:, 4]
sampled_box[:, 5] = sampled_dict["gt_boxes"][:, 3]
for i in range(sampled_box.shape[0]):
sampled_box[i,
6] = angle_in_limit(-sampled_dict["gt_boxes"][i,
6])
gt_box3d = np.concatenate([gt_box3d, sampled_box], axis=0)
classes = np.concatenate([classes, sampled_dict["gt_names"]])
lidar = np.concatenate([lidar, sampled_dict["points"]], axis=0)
if aug_pc:
choice = np.random.randint(0, 10)
if choice < 4:
# global rotation
angle = np.random.uniform(-np.pi / 30, np.pi / 30)
lidar[:, 0:3] = point_transform(lidar[:, 0:3], 0, 0, 0, rz=angle)
gt_box3d = box_transform(gt_box3d,
0,
0,
0,
r=angle,
coordinate='lidar')
newtag = 'aug_{}_1_{:.4f}'.format(tag, angle).replace('.', '_')
elif choice < 7:
# global translation
tx = np.random.uniform(-0.1, -0.1)
ty = np.random.uniform(-0.1, -0.1)
tz = np.random.uniform(-0.15, -0.15)
lidar[:, 0:3] = point_transform(lidar[:, 0:3], tx, ty, tz)
gt_box3d = box_transform(gt_box3d, tx, ty, tz, coordinate='lidar')
newtag = 'aug_{}_2_trans'.format(tag).replace('.', '_')
else:
# global scaling
factor = np.random.uniform(0.95, 1.05)
lidar[:, 0:3] = lidar[:, 0:3] * factor
gt_box3d[:, 0:6] = gt_box3d[:, 0:6] * factor
newtag = 'aug_{}_3_{:.4f}'.format(tag, factor).replace('.', '_')
else:
newtag = tag
P, Tr, R = load_calib(os.path.join(cfg.CALIB_DIR, tag + '.txt'))
label = box3d_to_label(gt_box3d[np.newaxis, ...],
classes[np.newaxis, ...],
coordinate='lidar',
P2=P,
T_VELO_2_CAM=Tr,
R_RECT_0=R)[0] # (N')
voxel_dict = process_pointcloud(tag, lidar)
if use_newtag:
return newtag, rgb, lidar, voxel_dict, label
else:
return tag, rgb, lidar, voxel_dict, label
fout.write('{} {} {}\n'.format(p[0], p[1], p[2]))
fout.close()
if __name__ == '__main__':
eae = eval_AE()
val_dir = os.path.join(cfg.DATA_DIR, 'training')
f_lidar = glob.glob(os.path.join(val_dir, 'velodyne', '*.bin'))
print('Total {} file'.format(len(f_lidar)))
c = np.random.randint(len(f_lidar))
print('choose {}'.format(f_lidar[c]))
lidar = np.fromfile(f_lidar[c], dtype=np.float32).reshape((-1, 4))
if cfg.FOV_FILTER:
lidar = extract_lidar_in_fov(lidar)
voxel_dict = process_pointcloud('pc', lidar)
num_voxels = voxel_dict['feature_buffer'].shape[0]
assert num_voxels == voxel_dict['coordinate_buffer'].shape[0]
print('num_voxels', num_voxels)
# generate mesh grid
linx = np.linspace(0, cfg.VOXEL_X_SIZE, cfg.VOXVOX_GRID_SIZE[0])
liny = np.linspace(0, cfg.VOXEL_Y_SIZE, cfg.VOXVOX_GRID_SIZE[1])
linz = np.linspace(0, cfg.VOXEL_Z_SIZE, cfg.VOXVOX_GRID_SIZE[2])
mesh = np.meshgrid(linx, liny, linz)
linx = np.expand_dims(mesh[0], axis=-1) # (4, 4, 8, 1)
liny = np.expand_dims(mesh[1], axis=-1)
linz = np.expand_dims(mesh[2], axis=-1)
mesh_coord = np.concatenate([linx, liny, linz], axis=-1)
print('mesh_coord', mesh_coord.shape)
def run(self):
raw_lidar = self.np_p_ranged
# print(raw_lidar.shape) # DEBUG
voxel, point_cloud = process_pointcloud(raw_lidar)
return raw_lidar, voxel, point_cloud
