def create_groundtruth_database(data_path,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True,
lidar_only=False,
bev_only=False,
coors_range=None):
RGB_embedding = True
root_path = pathlib.Path(data_path)
if info_path is None:
info_path = root_path / 'kitti_infos_train.pkl'
if database_save_path is None:
database_save_path = root_path / 'gt_database'
else:
database_save_path = pathlib.Path(database_save_path)
if db_info_save_path is None:
if RGB_embedding:
db_info_save_path = root_path / "kitti_dbinfos_train_RGB.pkl"
else:
db_info_save_path = root_path / "kitti_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
with open(info_path, 'rb') as f:
kitti_infos = pickle.load(f)
all_db_infos = {}
if used_classes is None:
used_classes = list(kitti.get_classes())
used_classes.pop(used_classes.index('DontCare'))
for name in used_classes:
all_db_infos[name] = []
group_counter = 0
for info in prog_bar(kitti_infos):
velodyne_path = info['velodyne_path']
if relative_path:
# velodyne_path = str(root_path / velodyne_path) + "_reduced"
velodyne_path = str(root_path / velodyne_path)
num_features = 4
if 'pointcloud_num_features' in info:
num_features = info['pointcloud_num_features']
points = np.fromfile(velodyne_path, dtype=np.float32,
count=-1).reshape([-1, num_features])
image_idx = info["image_idx"]
rect = info['calib/R0_rect']
P2 = info['calib/P2']
Trv2c = info['calib/Tr_velo_to_cam']
if not lidar_only:
points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2,
info["img_shape"])
annos = info["annos"]
names = annos["name"]
bboxes = annos["bbox"]
difficulty = annos["difficulty"]
gt_idxes = annos["index"]
num_obj = np.sum(annos["index"] >= 0)
rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj]
rbbox_lidar = box_np_ops.box_camera_to_lidar(rbbox_cam, rect, Trv2c)
if bev_only: # set z and h to limits
assert coors_range is not None
rbbox_lidar[:, 2] = coors_range[2]
rbbox_lidar[:, 5] = coors_range[5] - coors_range[2]
if RGB_embedding:
RGB_image = cv2.imread(str(root_path / info['img_path']))
points_camera = box_np_ops.box_lidar_to_camera(
points[:, :3], rect, Trv2c)
points_to_image_idx = box_np_ops.project_to_image(
points_camera, P2)
points_to_image_idx = points_to_image_idx.astype(int)
mask = box_np_ops.remove_points_outside_image(
RGB_image, points_to_image_idx)
points = points[mask]
points_to_image_idx = points_to_image_idx[mask]
BGR = RGB_image[points_to_image_idx[:, 1], points_to_image_idx[:,
0]]
points = np.concatenate((points, BGR), axis=1)
group_dict = {}
group_ids = np.full([bboxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(bboxes.shape[0], dtype=np.int64)
point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar)
for i in range(num_obj):
if RGB_embedding:
filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}_RGB.bin"
else:
filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}.bin"
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= rbbox_lidar[i, :3]
with open(filepath, 'w') as f:
gt_points.tofile(f)
if names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": gt_idxes[i],
"box3d_lidar": rbbox_lidar[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
# "group_id": -1,
# "bbox": bboxes[i],
}
local_group_id = group_ids[i]
# if local_group_id >= 0:
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
all_db_infos[names[i]].append(db_info)
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def create_groundtruth_database(data_path,
train_or_test,
info_path=None,
used_classes=None,
database_save_path=None,
db_info_save_path=None,
relative_path=True,
lidar_only=False,
bev_only=False,
coors_range=None):
# custom dataset parameter
custom_dataset = True # if true, it indicates that we are not operating on the kitti data
sample_val_dataset_mode = True if train_or_test == "test" else False # if true, we are creating a gt database for the test set (instead of the train set)
# ------------------------------------------------------------------------------------------------------
# create path where the gt boxes are stored
# -----------------------------------------------------------------------------------------------------
root_path = pathlib.Path(data_path)
if info_path is None:
if sample_val_dataset_mode:
info_path = root_path / 'kitti_infos_val.pkl'
else:
info_path = root_path / 'kitti_infos_train.pkl'
if database_save_path is None:
if sample_val_dataset_mode:
database_save_path = root_path / 'gt_database_val'
else:
database_save_path = root_path / 'gt_database'
else:
database_save_path = pathlib.Path(database_save_path)
if db_info_save_path is None:
if sample_val_dataset_mode:
db_info_save_path = root_path / "kitti_dbinfos_val.pkl"
else:
db_info_save_path = root_path / "kitti_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
# ------------------------------------------------------------------------------------------------------
# load kitti infos
# -----------------------------------------------------------------------------------------------------
with open(info_path, 'rb') as f:
kitti_infos = pickle.load(f)
all_db_infos = {}
# get the classnames we are intered in
if used_classes is None:
used_classes = list(kitti.get_classes())
used_classes.pop(used_classes.index('DontCare'))
for name in used_classes:
all_db_infos[name] = []
group_counter = 0
# ------------------------------------------------------------------------------------------------------
# iterate over kitti_infos
# -----------------------------------------------------------------------------------------------------
for info in prog_bar(kitti_infos):
# ------------------------------------------------------------------------------------------------------
# load pc
# -----------------------------------------------------------------------------------------------------
velodyne_path = info['velodyne_path']
if relative_path:
# velodyne_path = str(root_path / velodyne_path) + "_reduced"
velodyne_path = str(root_path / velodyne_path)
num_features = 4
if 'pointcloud_num_features' in info:
num_features = info['pointcloud_num_features']
if custom_dataset:
with open(str(velodyne_path)[:-3] + "pkl", 'rb') as file:
points = pickle.load(file, encoding='latin1')
else:
points = np.fromfile(str(velodyne_path),
dtype=np.float32,
count=-1).reshape([-1, 4])
image_idx = info["image_idx"]
rect = info['calib/R0_rect']
P2 = info['calib/P2']
Trv2c = info['calib/Tr_velo_to_cam']
# ------------------------------------------------------------------------------------------------------
# remove boxes outside the frustum of the image
# -----------------------------------------------------------------------------------------------------
if not lidar_only and not custom_dataset:
points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2,
info["img_shape"])
# ------------------------------------------------------------------------------------------------------
# get the bboxes and transform (annos not points)
# -----------------------------------------------------------------------------------------------------
annos = info["annos"]
names = annos["name"]
bboxes = annos["bbox"]
difficulty = annos["difficulty"]
gt_idxes = annos["index"]
num_obj = np.sum(annos["index"] >= 0)
rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj]
rbbox_lidar = box_np_ops.box_camera_to_lidar(rbbox_cam, rect, Trv2c)
if bev_only: # set z and h to limits
assert coors_range is not None
rbbox_lidar[:, 2] = coors_range[2]
rbbox_lidar[:, 5] = coors_range[5] - coors_range[2]
# ------------------------------------------------------------------------------------------------------
# other stuff
# -----------------------------------------------------------------------------------------------------
group_dict = {}
group_ids = np.full([bboxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(bboxes.shape[0], dtype=np.int64)
# ------------------------------------------------------------------------------------------------------
# test which points are in bboxes
# -----------------------------------------------------------------------------------------------------
point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar)
# ------------------------------------------------------------------------------------------------------
# iterate over all objects in the annos
# -----------------------------------------------------------------------------------------------------
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}.bin"
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= rbbox_lidar[i, :3]
# ------------------------------------------------------------------------------------------------------
# save points of gt boxes to files
# -----------------------------------------------------------------------------------------------------
with open(str(filepath)[:-3] + "pkl", 'wb') as file:
pickle.dump(np.array(gt_points), file, 2)
# debug
# with open("/home/makr/Documents/uni/TU/3.Master/experiments/own/tf_3dRGB_pc/debug_rviz/points/bbox_pixels.pkl", 'wb') as file:
# pickle.dump(np.array(gt_points), file, 2)
# ------------------------------------------------------------------------------------------------------
# save infos of gt boxes to single file
# -----------------------------------------------------------------------------------------------------
if names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": gt_idxes[i],
"box3d_lidar": rbbox_lidar[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
# "group_id": -1,
# "bbox": bboxes[i],
}
local_group_id = group_ids[i]
# if local_group_id >= 0:
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
all_db_infos[names[i]].append(db_info)
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def create_groundtruth_database(
data_path='/mrtstorage/datasets/kitti/object_detection',
info_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second'
'/kitti_infos_train.pkl',
used_classes=None,
database_save_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second'
'/gt_database',
db_info_save_path='/home/zhwang/second.pytorch/second/data/sets/kitti_second'
'/kitti_dbinfos_train.pkl',
relative_path=True,
lidar_only=False,
bev_only=False,
coors_range=None):
root_path = pathlib.Path(data_path)
if info_path is None:
info_path = root_path / 'kitti_infos_train.pkl'
if database_save_path is None:
database_save_path = root_path / 'gt_database'
else:
database_save_path = pathlib.Path(database_save_path)
if db_info_save_path is None:
db_info_save_path = root_path / "kitti_dbinfos_train.pkl"
database_save_path.mkdir(parents=True, exist_ok=True)
with open(info_path, 'rb') as f:
kitti_infos = pickle.load(f)
all_db_infos = {}
if used_classes is None:
used_classes = list(kitti.get_classes())
used_classes.pop(used_classes.index('DontCare'))
for name in used_classes:
all_db_infos[name] = []
group_counter = 0
for info in prog_bar(kitti_infos):
velodyne_path = info['velodyne_path']
if relative_path:
# velodyne_path = str(root_path / velodyne_path) + "_reduced"
velodyne_path = str(root_path / velodyne_path)
num_features = 4
if 'pointcloud_num_features' in info:
num_features = info['pointcloud_num_features']
points = np.fromfile(velodyne_path, dtype=np.float32,
count=-1).reshape([-1, num_features])
image_idx = info["image_idx"]
rect = info['calib/R0_rect']
P2 = info['calib/P2']
Trv2c = info['calib/Tr_velo_to_cam']
if not lidar_only:
points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2,
info["img_shape"])
"""
fusion camera data to lidar
"""
image_label = _get_semantic_segmentation_result(
image_idx,
data_dir='/home/zhwang/semantic-segmentation'
'/kitti_train_results')
points = _add_class_score(image_label, points, rect, Trv2c, P2)
annos = info["annos"]
names = annos["name"]
bboxes = annos["bbox"]
difficulty = annos["difficulty"]
gt_idxes = annos["index"]
num_obj = np.sum(annos["index"] >= 0)
rbbox_cam = kitti.anno_to_rbboxes(annos)[:num_obj]
rbbox_lidar = box_np_ops.box_camera_to_lidar(rbbox_cam, rect, Trv2c)
if bev_only: # set z and h to limits
assert coors_range is not None
rbbox_lidar[:, 2] = coors_range[2]
rbbox_lidar[:, 5] = coors_range[5] - coors_range[2]
group_dict = {}
group_ids = np.full([bboxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(bboxes.shape[0], dtype=np.int64)
point_indices = box_np_ops.points_in_rbbox(points, rbbox_lidar)
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{gt_idxes[i]}.bin"
filepath = database_save_path / filename
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= rbbox_lidar[i, :3]
with open(filepath, 'w') as f:
gt_points.tofile(f)
if names[i] in used_classes:
if relative_path:
db_path = str(database_save_path.stem + "/" + filename)
else:
db_path = str(filepath)
db_info = {
"name": names[i],
"path": db_path,
"image_idx": image_idx,
"gt_idx": gt_idxes[i],
"box3d_lidar": rbbox_lidar[i],
"num_points_in_gt": gt_points.shape[0],
"difficulty": difficulty[i],
# "group_id": -1,
# "bbox": bboxes[i],
}
local_group_id = group_ids[i]
# if local_group_id >= 0:
if local_group_id not in group_dict:
group_dict[local_group_id] = group_counter
group_counter += 1
db_info["group_id"] = group_dict[local_group_id]
if "score" in annos:
db_info["score"] = annos["score"][i]
all_db_infos[names[i]].append(db_info)
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(db_info_save_path, 'wb') as f:
pickle.dump(all_db_infos, f)
def create_inference_dataset(det_anno_path, calib_path, velodyne_path,
box3d_expansion, bbox_expansion, output_file):
# get annos
annos = get_label_annos(det_anno_path)
# create examples
examples = []
for anno in tqdm(annos):
# continue if zero object detected
if anno['image_idx'].shape[0] == 0:
continue
# get scene index
scene_idx = str(anno['image_idx'][0]).zfill(6)
# get calib
calib = get_kitti_calib(os.path.join(calib_path, scene_idx + '.txt'),
True)
# get box3d_camera
box3d_camera = anno_to_rbboxes(anno)
box3d_lidar = box_camera_to_lidar(box3d_camera, calib["R0_rect"],
calib["Tr_velo_to_cam"])
box3d_lidar[:, 2] += box3d_lidar[:, 5] / 2
# get expanded box3d
box3d_lidar_expanded = expand_box3d(box3d_lidar, box3d_expansion)
# get bbox
bbox = box3d_to_bbox(box3d_lidar_expanded, calib["R0_rect"],
calib["Tr_velo_to_cam"], calib['P2'])
bbox_expanded = expand_bbox(bbox, bbox_expansion).astype(np.int)
bbox_expanded[:, 0] = np.clip(bbox_expanded[:, 0], 0, 1242)
bbox_expanded[:, 1] = np.clip(bbox_expanded[:, 1], 0, 375)
bbox_expanded[:, 2] = np.clip(bbox_expanded[:, 2], 0, 1242)
bbox_expanded[:, 3] = np.clip(bbox_expanded[:, 3], 0, 375)
# read scene pts
pts = read_bin(os.path.join(velodyne_path, scene_idx + '.bin'))[:, :3]
# create example
for idx in range(box3d_lidar_expanded.shape[0]):
filtered_pts = points_in_rbbox(
pts, box3d_lidar_expanded[idx][np.newaxis, ...])
res = {
'rgb': {},
'point': {},
'calib': calib,
}
res['rgb']['bbox'] = bbox_expanded[idx]
res['rgb']['img_id'] = anno['image_idx'][0]
res['point']['box3d_lidar'] = box3d_lidar_expanded[idx]
res['point']['point_dict'] = {
'source': [pts[filtered_pts.reshape(-1)]],
'gt': [],
3: [],
4: [],
5: [],
7: [],
9: [],
}
examples.append(res)
# write to file
with open(output_file, 'wb') as f:
pickle.dump(examples, f)
