def box3d_to_bbox(box3d, rect, Trv2c, P2):
box3d_lidar = box3d.copy()
box3d_lidar[:, 2] -= box3d_lidar[:, 5] / 2
box3d_camera = box_lidar_to_camera(box3d_lidar, rect, Trv2c)
box_corners = center_to_corner_box3d(box3d_camera[:, :3],
box3d_camera[:, 3:6],
box3d_camera[:, 6], [0.5, 1.0, 0.5],
axis=1)
box_corners_in_image = project_to_image(box_corners, P2)
# box_corners_in_image: [N, 8, 2]
minxy = np.min(box_corners_in_image, axis=1)
maxxy = np.max(box_corners_in_image, axis=1)
bbox = np.concatenate([minxy, maxxy], axis=1)
return bbox
def get_box3d_cam(box3d_lidar, rect, Trv2c): from second.core.box_np_ops import box_lidar_to_camera box3d_cam = box_lidar_to_camera(box3d_lidar, rect, Trv2c) return box3d_cam
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 convert_detection_to_kitti_annos(self, detection):
class_names = self._class_names
det_image_idxes = [det["metadata"]["image_idx"] for det in detection]
gt_image_idxes = [
info["image"]["image_idx"] for info in self._kitti_infos
]
annos = []
for i in range(len(detection)):
det_idx = det_image_idxes[i]
det = detection[i]
# info = self._kitti_infos[gt_image_idxes.index(det_idx)]
info = self._kitti_infos[i]
calib = info["calib"]
rect = calib["R0_rect"]
Trv2c = calib["Tr_velo_to_cam"]
P2 = calib["P2"]
final_box_preds = det["box3d_lidar"].detach().cpu().numpy()
label_preds = det["label_preds"].detach().cpu().numpy()
scores = det["scores"].detach().cpu().numpy()
if final_box_preds.shape[0] != 0:
final_box_preds[:, 2] -= final_box_preds[:, 5] / 2
box3d_camera = box_np_ops.box_lidar_to_camera(
final_box_preds, rect, Trv2c)
locs = box3d_camera[:, :3]
dims = box3d_camera[:, 3:6]
angles = box3d_camera[:, 6]
camera_box_origin = [0.5, 1.0, 0.5]
box_corners = box_np_ops.center_to_corner_box3d(
locs, dims, angles, camera_box_origin, axis=1)
box_corners_in_image = box_np_ops.project_to_image(
box_corners, P2)
# box_corners_in_image: [N, 8, 2]
minxy = np.min(box_corners_in_image, axis=1)
maxxy = np.max(box_corners_in_image, axis=1)
bbox = np.concatenate([minxy, maxxy], axis=1)
anno = kitti.get_start_result_anno()
num_example = 0
box3d_lidar = final_box_preds
for j in range(box3d_lidar.shape[0]):
image_shape = info["image"]["image_shape"]
if bbox[j, 0] > image_shape[1] or bbox[j, 1] > image_shape[0]:
continue
if bbox[j, 2] < 0 or bbox[j, 3] < 0:
continue
bbox[j, 2:] = np.minimum(bbox[j, 2:], image_shape[::-1])
bbox[j, :2] = np.maximum(bbox[j, :2], [0, 0])
anno["bbox"].append(bbox[j])
# convert center format to kitti format
# box3d_lidar[j, 2] -= box3d_lidar[j, 5] / 2
anno["alpha"].append(
-np.arctan2(-box3d_lidar[j, 1], box3d_lidar[j, 0]) +
box3d_camera[j, 6])
anno["dimensions"].append(box3d_camera[j, 3:6])
anno["location"].append(box3d_camera[j, :3])
anno["rotation_y"].append(box3d_camera[j, 6])
anno["name"].append(class_names[int(label_preds[j])])
anno["truncated"].append(0.0)
anno["occluded"].append(0)
anno["score"].append(scores[j])
num_example += 1
if num_example != 0:
anno = {n: np.stack(v) for n, v in anno.items()}
annos.append(anno)
else:
annos.append(kitti.empty_result_anno())
num_example = annos[-1]["name"].shape[0]
annos[-1]["metadata"] = det["metadata"]
return annos
def prep_pointcloud(
input_dict,
root_path,
# voxel_generator,
fv_generator,
target_assigner,
db_sampler=None,
max_voxels=20000,
class_names=['Car'],
remove_outside_points=False,
training=True,
create_targets=True,
shuffle_points=False,
reduce_valid_area=False,
remove_unknown=False,
gt_rotation_noise=[-np.pi / 3, np.pi / 3],
gt_loc_noise_std=[1.0, 1.0, 1.0],
global_rotation_noise=[-np.pi / 4, np.pi / 4],
global_scaling_noise=[0.95, 1.05],
global_loc_noise_std=(0.2, 0.2, 0.2),
global_random_rot_range=[0.78, 2.35],
generate_bev=False,
without_reflectivity=False,
num_point_features=4,
anchor_area_threshold=1,
gt_points_drop=0.0,
gt_drop_max_keep=10,
remove_points_after_sample=False,
anchor_cache=None,
remove_environment=False,
random_crop=False,
reference_detections=None,
add_rgb_to_points=False,
lidar_input=False,
unlabeled_db_sampler=None,
out_size_factor=2,
min_gt_point_dict=None,
bev_only=False,
use_group_id=False,
out_dtype=np.float32,
num_classes=1,
RGB_embedding=False):
"""convert point cloud to voxels, create targets if ground truths
exists.
"""
# prep_pointcloud_start = time.time()
points = input_dict["points"]
# if training:
gt_boxes = input_dict["gt_boxes"]
gt_names = input_dict["gt_names"]
difficulty = input_dict["difficulty"]
group_ids = None
if use_group_id and "group_ids" in input_dict:
group_ids = input_dict["group_ids"]
rect = input_dict["rect"]
Trv2c = input_dict["Trv2c"]
P2 = input_dict["P2"]
unlabeled_training = unlabeled_db_sampler is not None
image_idx = input_dict["image_idx"]
# t1 = time.time() - prep_pointcloud_start
if shuffle_points:
# shuffle is a little slow.
np.random.shuffle(points)
# t2 = time.time() - prep_pointcloud_start
# print("t2-t1: ", t2-t1) # 0.035
if reference_detections is not None:
C, R, T = box_np_ops.projection_matrix_to_CRT_kitti(P2)
frustums = box_np_ops.get_frustum_v2(reference_detections, C)
frustums -= T
# frustums = np.linalg.inv(R) @ frustums.T
frustums = np.einsum('ij, akj->aki', np.linalg.inv(R), frustums)
frustums = box_np_ops.camera_to_lidar(frustums, rect, Trv2c)
surfaces = box_np_ops.corner_to_surfaces_3d_jit(frustums)
masks = points_in_convex_polygon_3d_jit(points, surfaces)
points = points[masks.any(-1)]
if remove_outside_points: # and not lidar_input:
image_shape = input_dict["image_shape"]
points = box_np_ops.remove_outside_points(points, rect, Trv2c, P2,
image_shape)
if remove_environment is True: # and training:
selected = kitti.keep_arrays_by_name(gt_names, class_names)
gt_boxes = gt_boxes[selected]
gt_names = gt_names[selected]
difficulty = difficulty[selected]
if group_ids is not None:
group_ids = group_ids[selected]
points = prep.remove_points_outside_boxes(points, gt_boxes)
# if training:
# print(gt_names)
selected = kitti.drop_arrays_by_name(gt_names, ["DontCare"])
gt_boxes = gt_boxes[selected]
gt_names = gt_names[selected]
difficulty = difficulty[selected]
if group_ids is not None:
group_ids = group_ids[selected]
# t3 = time.time() - prep_pointcloud_start
# print("t3-t2: ", t3 - t2) # 0.0002
gt_boxes = box_np_ops.box_camera_to_lidar(gt_boxes, rect, Trv2c)
if remove_unknown:
remove_mask = difficulty == -1
"""
gt_boxes_remove = gt_boxes[remove_mask]
gt_boxes_remove[:, 3:6] += 0.25
points = prep.remove_points_in_boxes(points, gt_boxes_remove)
"""
keep_mask = np.logical_not(remove_mask)
gt_boxes = gt_boxes[keep_mask]
gt_names = gt_names[keep_mask]
difficulty = difficulty[keep_mask]
if group_ids is not None:
group_ids = group_ids[keep_mask]
gt_boxes_mask = np.array([n in class_names for n in gt_names],
dtype=np.bool_)
# t4 = time.time() - prep_pointcloud_start
# print("t4-t3: ", t4 - t3) # 0.001
if RGB_embedding:
RGB_image = cv2.imread(input_dict['image_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)
# test_mask = points_camera[mask][:, 0] < 0
# test_image_idx = points_to_image_idx[test_mask]
# RGB_image[test_image_idx[:, 1], test_image_idx[:, 0]] = [255, 0, 0]
# test_mask = points_camera[mask][:, 0] >= 0
# test_image_idx = points_to_image_idx[test_mask]
# RGB_image[test_image_idx[:, 1], test_image_idx[:, 0]] = [0, 0, 255]
# print()
# t5 = time.time() - prep_pointcloud_start
# print("t5-t4: ", t5 - t4) # 0.019
# TODO
if db_sampler is not None and training: # and not RGB_embedding:
if RGB_embedding:
num_point_features += 3
fg_points_mask = box_np_ops.points_in_rbbox(points, gt_boxes)
fg_points_list = []
bg_points_mask = np.zeros((points.shape[0]), dtype=bool)
for i in range(fg_points_mask.shape[1]):
fg_points_list.append(points[fg_points_mask[:, i]])
bg_points_mask = np.logical_or(bg_points_mask, fg_points_mask[:,
i])
bg_points_mask = np.logical_not(bg_points_mask)
sampled_dict = db_sampler.sample_all(root_path,
points[bg_points_mask],
gt_boxes,
gt_names,
fg_points_list,
num_point_features,
random_crop,
gt_group_ids=group_ids,
rect=rect,
Trv2c=Trv2c,
P2=P2)
# sampled_dict = db_sampler.sample_all(
# root_path,
# gt_boxes,
# gt_names,
# num_point_features,
# random_crop,
# gt_group_ids=group_ids,
# rect=rect,
# Trv2c=Trv2c,
# P2=P2)
# t_sample_all = time.time() - prep_pointcloud_start
# print("t_sample_all - t5: ", t_sample_all - t5) # 3.83
if sampled_dict is not None:
sampled_gt_names = sampled_dict["gt_names"]
sampled_gt_boxes = sampled_dict["gt_boxes"]
points = sampled_dict["points"]
sampled_gt_masks = sampled_dict["gt_masks"]
remained_boxes_idx = sampled_dict["remained_boxes_idx"]
# gt_names = gt_names[gt_boxes_mask].tolist()
gt_names = np.concatenate([gt_names, sampled_gt_names], axis=0)
# gt_names += [s["name"] for s in sampled]
gt_boxes = np.concatenate([gt_boxes, sampled_gt_boxes])
gt_boxes_mask = np.concatenate([gt_boxes_mask, sampled_gt_masks],
axis=0)
gt_names = gt_names[remained_boxes_idx]
gt_boxes = gt_boxes[remained_boxes_idx]
gt_boxes_mask = gt_boxes_mask[remained_boxes_idx]
if group_ids is not None:
sampled_group_ids = sampled_dict["group_ids"]
group_ids = np.concatenate([group_ids, sampled_group_ids])
group_ids = group_ids[remained_boxes_idx]
# if remove_points_after_sample:
# # points = prep.remove_points_in_boxes(
# # points, sampled_gt_boxes)
# locs = sampled_gt_boxes[:, 0:3]
# dims = sampled_gt_boxes[:, 3:6]
# angles = sampled_gt_boxes[:, 6]
# camera_box_origin = [0.5, 0.5, 0]
#
# box_corners = box_np_ops.center_to_corner_box3d(
# locs, dims, angles, camera_box_origin, axis=2)
# box_corners_in_image = box_np_ops.project_to_fv_image(
# box_corners, example['grid_size'][i], example['meta'][i])
# box_centers_in_image = box_np_ops.project_to_fv_image(
# locs, example['grid_size'][i], example['meta'][i])
# t_sample_all2 = time.time() - prep_pointcloud_start
# print("t_sample_all2 - t_sample_all: ", t_sample_all2 - t_sample_all) # 0.0002
# unlabeled_mask = np.zeros((gt_boxes.shape[0], ), dtype=np.bool_)
# if without_reflectivity and training:
# used_point_axes = list(range(num_point_features))
# used_point_axes.pop(3)
# points = points[:, used_point_axes]
# pc_range = voxel_generator.point_cloud_range
# bev_only = False
# if bev_only: # set z and h to limits
# gt_boxes[:, 2] = pc_range[2]
# gt_boxes[:, 5] = pc_range[5] - pc_range[2]
if training:
gt_loc_noise_std = [0.0, 0.0, 0.0]
prep.noise_per_object_v3_(
gt_boxes,
points,
gt_boxes_mask,
rotation_perturb=gt_rotation_noise,
center_noise_std=gt_loc_noise_std,
global_random_rot_range=global_random_rot_range,
group_ids=group_ids,
num_try=100)
# t_noise = time.time() - prep_pointcloud_start
# print("t_noise - t_sample_all2: ", t_noise - t_sample_all2) # 12.01
# should remove unrelated objects after noise per object
gt_boxes = gt_boxes[gt_boxes_mask]
gt_names = gt_names[gt_boxes_mask]
if group_ids is not None:
group_ids = group_ids[gt_boxes_mask]
gt_classes = np.array([class_names.index(n) + 1 for n in gt_names],
dtype=np.int32)
# t6 = time.time() - prep_pointcloud_start
# print("t6-t5: ", t6 - t5) # 16.0
if training:
gt_boxes, points = prep.random_flip(gt_boxes, points)
# gt_boxes, points = prep.global_rotation(
# gt_boxes, points, rotation=global_rotation_noise)
gt_boxes, points = prep.global_scaling_v2(gt_boxes, points,
*global_scaling_noise)
# Global translation
# gt_boxes, points = prep.global_translate(gt_boxes, points, global_loc_noise_std)
# bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
bv_range = [0, -40, 70.4, 40]
mask = prep.filter_gt_box_outside_range(gt_boxes, bv_range)
gt_boxes = gt_boxes[mask]
gt_classes = gt_classes[mask]
if group_ids is not None:
group_ids = group_ids[mask]
# limit rad to [-pi, pi]
gt_boxes[:, 6] = box_np_ops.limit_period(gt_boxes[:, 6],
offset=0.5,
period=2 * np.pi)
# TODO
# if shuffle_points:
# # shuffle is a little slow.
# np.random.shuffle(points)
# # t7 = time.time() - prep_pointcloud_start
# # print("t7-t6: ", t7 - t6) # 1.95
# voxels, coordinates, num_points = voxel_generator.generate(
# points, max_voxels, RGB_embedding=RGB_embedding)
# # t8 = time.time() - prep_pointcloud_start
# # print("t8-t7: ", t8 - t7) # 2.0
# voxel_size = voxel_generator.voxel_size
# grid_size = voxel_generator.grid_size
# pc_range = copy.deepcopy(voxel_generator.point_cloud_range)
# grid_size = voxel_generator.grid_size
# phi_min = voxel_generator.phi_min
# theta_min = voxel_generator.theta_min
# image_h, image_w = grid_size[1], grid_size[0]
# c = np.array([image_w / 2., image_h / 2.])
# s = np.array([image_w, image_h], dtype=np.int32)
# meta = {'c': c, 's': s, 'calib': P2, 'phi_min': phi_min, 'theta_min': theta_min}
# t7 = time.time() - prep_pointcloud_start
# print("t7-t6: ", t7 - t6) # 1.95
fv_image, points_mask = fv_generator.generate(points,
RGB_embedding=RGB_embedding,
occupancy_embedding=False)
# t8 = time.time() - prep_pointcloud_start
# print("t8-t7: ", t8 - t7) # 2.0
fv_dim = fv_generator.fv_dim
pc_range = copy.deepcopy(fv_generator.spherical_coord_range)
grid_size = fv_generator.grid_size
phi_min = fv_generator.phi_min
theta_min = fv_generator.theta_min
image_h, image_w = fv_dim[1], fv_dim[0]
c = np.array([image_w / 2., image_h / 2.])
s = np.array([image_w, image_h], dtype=np.int32)
meta = {
'c': c,
's': s,
'calib': P2,
'phi_min': phi_min,
'theta_min': theta_min
}
fv_image = np.transpose(fv_image, [2, 1, 0])
max_objs = 50
num_objs = min(gt_boxes.shape[0], max_objs)
box_np_ops.change_box3d_center_(gt_boxes,
src=[0.5, 0.5, 0],
dst=[0.5, 0.5, 0.5])
spherical_gt_boxes = np.zeros((max_objs, gt_boxes.shape[1]))
spherical_gt_boxes[:num_objs, :] = gt_boxes[:num_objs, :]
spherical_gt_boxes[:num_objs, :] = convert_to_spherical_coor(
gt_boxes[:num_objs, :])
spherical_gt_boxes[:num_objs, 0] -= phi_min
spherical_gt_boxes[:num_objs, 1] -= theta_min
spherical_gt_boxes[:num_objs, 0] /= grid_size[0]
spherical_gt_boxes[:num_objs, 1] /= grid_size[1]
spherical_gt_boxes, num_objs = filter_outside_range(
spherical_gt_boxes, num_objs, fv_dim)
# t9 = time.time() - prep_pointcloud_start
# print("t9-t8: ", t9 - t8) # 0.0005
example = {
'fv_image': fv_image,
'grid_size': grid_size,
'pc_range': pc_range,
'meta': meta,
'spherical_gt_boxes': spherical_gt_boxes,
'resized_image_shape': grid_size
}
example.update({'rect': rect, 'Trv2c': Trv2c, 'P2': P2})
if RGB_embedding:
RGB_image = cv2.resize(RGB_image, (image_w, image_h))
example.update({'RGB_image': RGB_image})
if training:
hm = np.zeros((num_classes, image_h, image_w), dtype=np.float32)
reg = np.zeros((max_objs, 2), dtype=np.float32)
dep = np.zeros((max_objs, 1), dtype=np.float32)
rotbin = np.zeros((max_objs, 2), dtype=np.int64)
rotres = np.zeros((max_objs, 2), dtype=np.float32)
dim = np.zeros((max_objs, 3), dtype=np.float32)
ind = np.zeros((max_objs), dtype=np.int64)
reg_mask = np.zeros((max_objs), dtype=np.uint8)
rot_mask = np.zeros((max_objs), dtype=np.uint8)
#
# hm = np.zeros((num_classes, image_h, image_w), dtype=np.float32)
# reg = np.zeros((image_w, image_h, 2), dtype=np.float32)
# dep = np.zeros((image_w, image_h, 1), dtype=np.float32)
# rotbin = np.zeros((image_w, image_h, 2), dtype=np.int64)
# rotres = np.zeros((image_w, image_h, 2), dtype=np.float32)
# dim = np.zeros((image_w, image_h, 3), dtype=np.float32)
# # ind = np.zeros((max_objs), dtype=np.int64)
# fg_mask = np.zeros((image_w, image_h), dtype=np.uint8)
# # rot_mask = np.zeros((max_objs), dtype=np.uint8)
draw_gaussian = draw_umich_gaussian
# center heatmap
radius = int(image_h / 30)
# radius = int(image_h / 25)
for k in range(num_objs):
gt_3d_box = spherical_gt_boxes[k]
cls_id = 0
# print('heatmap gaussian radius: ' + str(radius))
ct = np.array([gt_3d_box[0], gt_3d_box[1]], dtype=np.float32)
ct_int = ct.astype(np.int32)
draw_gaussian(hm[cls_id], ct, radius)
# depth(distance), wlh
dep[k] = gt_3d_box[2]
dim[k] = gt_3d_box[3:6]
# dep[ct_int[0], ct_int[1], 0] = gt_3d_box[2]
# dim[ct_int[0], ct_int[1], :] = gt_3d_box[3:6]
# reg, ind, mask
reg[k] = ct - ct_int
ind[k] = ct_int[1] * image_w + ct_int[0]
reg_mask[k] = rot_mask[k] = 1
# fg_mask[ct_int[0], ct_int[1]] = 1
# ry
ry = gt_3d_box[6]
if ry < np.pi / 6. or ry > 5 * np.pi / 6.:
rotbin[k, 0] = 1
rotres[k, 0] = ry - (-0.5 * np.pi)
# rotbin[ct_int[0], ct_int[1], 0] = 1
# rotres[ct_int[0], ct_int[1], 0] = ry - (-0.5 * np.pi)
if ry > -np.pi / 6. or ry < -5 * np.pi / 6.:
rotbin[k, 1] = 1
rotres[k, 1] = ry - (0.5 * np.pi)
# rotbin[ct_int[0], ct_int[1], 1] = 1
# rotres[ct_int[0], ct_int[1], 1] = ry - (0.5 * np.pi)
example.update({
'hm': hm,
'dep': dep,
'dim': dim,
'ind': ind,
'rotbin': rotbin,
'rotres': rotres,
'reg_mask': reg_mask,
'rot_mask': rot_mask,
'reg': reg
})
# example.update({
# 'hm': hm, 'dep': dep, 'dim': dim,
# 'rotbin': rotbin, 'rotres': rotres,
# 'fg_mask': fg_mask, 'reg': reg
# })
# t10 = time.time() - prep_pointcloud_start
# print("total: ", t10) # 19.58
return example
