def _calculate_num_points_in_gt(data_path,
infos,
relative_path,
remove_outside=False,
num_features=3):
'''
计算bounding box内点的数量
'''
for info in infos:
pc_info = info["point_cloud"]
if relative_path:
v_path = str(Path(data_path) / pc_info["velodyne_path"])
else:
v_path = pc_info["velodyne_path"]
pcd = o3d.io.read_point_cloud(str(v_path))
points_v = np.asarray(pcd.points)
if num_features >= 4:
normals_v = np.asarray(pcd.normals)
points_v = np.concatenate([points_v, normals_v],
axis=1)[:, :num_features]
annos = info["annos"]
dims = annos["dimensions"]
loc = annos["location"]
rots = annos["rotation_y"]
gt_boxes = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1)
indices = box_np_ops.points_in_rbbox(points_v[:, :3], gt_boxes)
num_points_in_gt = indices.sum(0)
annos["num_points_in_gt"] = num_points_in_gt
def _calculate_num_points_in_gt(data_path,
infos,
relative_path,
remove_outside=True,
num_features=4):
for info in infos:
pc_info = info["point_cloud"]
image_info = info["image"]
calib = info["calib"]
if relative_path:
v_path = str(Path(data_path) / pc_info["velodyne_path"])
else:
v_path = pc_info["velodyne_path"]
points_v = np.fromfile(v_path, dtype=np.float32,
count=-1).reshape([-1, num_features])
rect = calib["R0_rect"]
Trv2c = calib["Tr_velo_to_cam"]
P2 = calib["P2"]
if remove_outside:
points_v = box_np_ops.remove_outside_points(
points_v, rect, Trv2c, P2, image_info["image_shape"])
# points_v = points_v[points_v[:, 0] > 0]
annos = info["annos"]
num_obj = len([n for n in annos["name"] if n != "DontCare"])
# annos = kitti.filter_kitti_anno(annos, ['DontCare'])
dims = annos["dimensions"][:num_obj]
loc = annos["location"][:num_obj]
rots = annos["rotation_y"][:num_obj]
gt_boxes_camera = np.concatenate([loc, dims, rots[..., np.newaxis]],
axis=1)
gt_boxes_lidar = box_np_ops.box_camera_to_lidar(
gt_boxes_camera, rect, Trv2c)
indices = box_np_ops.points_in_rbbox(points_v[:, :3], gt_boxes_lidar)
num_points_in_gt = indices.sum(0)
num_ignored = len(annos["dimensions"]) - num_obj
num_points_in_gt = np.concatenate(
[num_points_in_gt, -np.ones([num_ignored])])
annos["num_points_in_gt"] = num_points_in_gt.astype(np.int32)
def create_groundtruth_database(
dataset_class_name,
data_path,
info_path=None,
used_classes=None,
db_path=None,
dbinfo_path=None,
relative_path=True,
add_rgb=False,
lidar_only=False,
bev_only=False,
coors_range=None,
**kwargs,
):
pipeline = [
{
"type": "LoadPointCloudFromFile",
"dataset": dataset_name_map[dataset_class_name],
},
{
"type": "LoadPointCloudAnnotations",
"with_bbox": True
},
]
if "nsweeps" in kwargs:
dataset = get_dataset(dataset_class_name)(
info_path=info_path,
root_path=data_path,
pipeline=pipeline,
test_mode=True,
nsweeps=kwargs["nsweeps"],
)
nsweeps = dataset.nsweeps
else:
dataset = get_dataset(dataset_class_name)(info_path=info_path,
root_path=data_path,
test_mode=True,
pipeline=pipeline)
nsweeps = 1
root_path = Path(data_path)
if dataset_class_name == "NUSC":
if db_path is None:
db_path = root_path / f"gt_database_{nsweeps}sweeps_withvelo"
if dbinfo_path is None:
dbinfo_path = root_path / f"dbinfos_train_{nsweeps}sweeps_withvelo.pkl"
else:
if db_path is None:
db_path = root_path / "gt_database"
if dbinfo_path is None:
dbinfo_path = root_path / "dbinfos_train.pkl"
if dataset_class_name == "NUSC" or dataset_class_name == "LYFT":
point_features = 5
elif dataset_class_name == "KITTI":
point_features = 4
db_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
# def prepare_single_data(index):
for index in tqdm(range(len(dataset))):
image_idx = index
# modified to nuscenes
sensor_data = dataset.get_sensor_data(index)
# for nsweep, sensor_data in enumerate(sensor_datas):
if "image_idx" in sensor_data["metadata"]:
image_idx = sensor_data["metadata"]["image_idx"]
if dataset_class_name == "NUSC":
points = sensor_data["lidar"]["combined"]
elif dataset_class_name == "KITTI":
points = sensor_data["lidar"]["points"]
elif dataset_class_name == "LYFT":
points = sensor_data["lidar"]["points"]
print(points.shape)
annos = sensor_data["lidar"]["annotations"]
gt_boxes = annos["boxes"]
names = annos["names"]
group_dict = {}
group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
if "difficulty" in annos:
difficulty = annos["difficulty"]
num_obj = gt_boxes.shape[0]
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes)
for i in range(num_obj):
filename = f"{image_idx}_{names[i]}_{i}.bin"
filepath = db_path / filename
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= gt_boxes[i, :3]
with open(filepath, "w") as f:
gt_points[:, :point_features].tofile(f)
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_dump_path = str(db_path.stem + "/" + filename)
else:
db_dump_path = str(filepath)
db_info = {
"name": names[i],
"path": db_dump_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[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]
if names[i] in all_db_infos:
all_db_infos[names[i]].append(db_info)
else:
all_db_infos[names[i]] = [db_info]
# print(f"Finish {index}th sample")
print("dataset length: ", len(dataset))
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(dbinfo_path, "wb") as f:
pickle.dump(all_db_infos, f)
def create_groundtruth_database(
dataset_class_name,
data_path,
info_path=None,
used_classes=None,
db_path=None,
dbinfo_path=None,
relative_path=True,
add_rgb=False,
lidar_only=False,
bev_only=False,
coors_range=None,
gt_aug_with_context=-1.0,
**kwargs,
):
gt_aug_with_context = gt_aug_with_context
pipeline = [
{
"type": "LoadPointCloudFromFile",
"dataset": dataset_name_map[dataset_class_name],
},
{
"type": "LoadPointCloudAnnotations",
"with_bbox": True,
"enable_difficulty_level": True
},
]
# get KittiDataset loaded points and annos.
dataset = get_dataset(dataset_class_name)(info_path=info_path,
root_path=data_path,
test_mode=True,
pipeline=pipeline)
# prepare dbinfo_path and db_path.
root_path = Path(data_path)
if db_path is None:
db_path = root_path / "gt_database"
if dbinfo_path is None:
dbinfo_path = root_path / "dbinfos_train.pkl"
db_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
for index in tqdm(range(len(dataset))):
image_idx = index
sensor_data = dataset.get_sensor_data(
index) # see in loading.py after pipelines.
if "image_idx" in sensor_data[
"metadata"]: # True, image_idx = file_name (000001)
image_idx = sensor_data["metadata"]["image_idx"]
points = sensor_data["lidar"]["points"]
annos = sensor_data["lidar"]["annotations"]
gt_boxes = annos[
"boxes"] # gt_boxes of all classes, dc has been removed in pipeline LoadPointCloudAnnotations.
names = annos["names"] # gt_names.
group_dict = {}
group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos: # False
group_ids = annos["group_ids"]
else:
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
if "difficulty" in annos: # False
difficulty = annos["difficulty"]
num_obj = gt_boxes.shape[0]
# todo: maybe we need add some contexual points here.
offset = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] # [x, y, z, w, l, h, ry]
if gt_aug_with_context > 0.0:
offset = [
0.0, 0.0, 0.0, gt_aug_with_context, gt_aug_with_context, 0.0,
0.0
]
db_path = root_path / "gt_enlarged_database"
dbinfo_path = root_path / "dbinfos_enlarged_train.pkl"
db_path.mkdir(parents=True, exist_ok=True)
point_indices_for_num = box_np_ops.points_in_rbbox(points, gt_boxes)
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes + offset)
for i in range(num_obj): # in a single scene.
filename = f"{image_idx}_{names[i]}_{i}.bin"
filepath = db_path / filename
gt_points = points[point_indices[:, i]]
num_points_in_gt = point_indices_for_num[:, i].sum()
gt_points[:, :3] -= gt_boxes[i, :
3] # only record relative distance
with open(filepath,
"w") as f: # db: gt points in each gt_box are saved
gt_points[:, :4].tofile(f)
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_dump_path = str(db_path.stem + "/" + filename)
else:
db_dump_path = str(filepath)
db_info = {
"name": names[i],
"path": db_dump_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[i],
"num_points_in_gt": num_points_in_gt,
"difficulty":
difficulty[i] # todo: not accurate, all are set as 0.
}
local_group_id = group_ids[i]
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] # count from 0 to total_num_of_specific_class[like 13442 for car]
if "score" in annos: # False
db_info["score"] = annos["score"][i]
if names[
i] in all_db_infos: # all_db_infos are grouped by class_names (like car, pedestrian, cyclist, cycle)
all_db_infos[names[i]].append(
db_info) # all db infos include info of all db
else:
all_db_infos[names[i]] = [db_info]
print("dataset length: ", len(dataset))
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(dbinfo_path, "wb") as f:
pickle.dump(all_db_infos, f)
def create_groundtruth_database(
dataset_class_name,
data_path,
info_path=None,
used_classes=None,
db_path=None,
dbinfo_path=None,
relative_path=True,
**kwargs,
):
pipeline = [
{
"type": "LoadPointCloudFromFile",
"dataset": dataset_name_map[dataset_class_name],
},
{
"type": "LoadPointCloudAnnotations",
"with_bbox": True
},
]
if "nsweeps" in kwargs:
dataset = get_dataset(dataset_class_name)(
info_path=info_path,
root_path=data_path,
pipeline=pipeline,
test_mode=True,
nsweeps=kwargs["nsweeps"],
)
nsweeps = dataset.nsweeps
else:
dataset = get_dataset(dataset_class_name)(info_path=info_path,
root_path=data_path,
test_mode=True,
pipeline=pipeline)
nsweeps = 1
root_path = Path(data_path)
if dataset_class_name in ["WAYMO", "NUSC"]:
if db_path is None:
db_path = root_path / f"gt_database_{nsweeps}sweeps_withvelo"
if dbinfo_path is None:
dbinfo_path = root_path / f"dbinfos_train_{nsweeps}sweeps_withvelo.pkl"
else:
raise NotImplementedError()
if dataset_class_name == "NUSC":
point_features = 5
elif dataset_class_name == "WAYMO":
point_features = 5 if nsweeps == 1 else 6
else:
raise NotImplementedError()
db_path.mkdir(parents=True, exist_ok=True)
all_db_infos = {}
group_counter = 0
for index in tqdm(range(len(dataset))):
image_idx = index
# modified to nuscenes
sensor_data = dataset.get_sensor_data(index)
if "image_idx" in sensor_data["metadata"]:
image_idx = sensor_data["metadata"]["image_idx"]
if nsweeps > 1:
points = sensor_data["lidar"]["combined"]
else:
points = sensor_data["lidar"]["points"]
annos = sensor_data["lidar"]["annotations"]
gt_boxes = annos["boxes"]
names = annos["names"]
if dataset_class_name == 'WAYMO':
# waymo dataset contains millions of objects and it is not possible to store
# all of them into a single folder
# we randomly sample a few objects for gt augmentation
# We keep all cyclist as they are rare
if index % 4 != 0:
mask = (names == 'VEHICLE')
mask = np.logical_not(mask)
names = names[mask]
gt_boxes = gt_boxes[mask]
if index % 2 != 0:
mask = (names == 'PEDESTRIAN')
mask = np.logical_not(mask)
names = names[mask]
gt_boxes = gt_boxes[mask]
group_dict = {}
group_ids = np.full([gt_boxes.shape[0]], -1, dtype=np.int64)
if "group_ids" in annos:
group_ids = annos["group_ids"]
else:
group_ids = np.arange(gt_boxes.shape[0], dtype=np.int64)
difficulty = np.zeros(gt_boxes.shape[0], dtype=np.int32)
if "difficulty" in annos:
difficulty = annos["difficulty"]
num_obj = gt_boxes.shape[0]
if num_obj == 0:
continue
point_indices = box_np_ops.points_in_rbbox(points, gt_boxes)
for i in range(num_obj):
if (used_classes is None) or names[i] in used_classes:
filename = f"{image_idx}_{names[i]}_{i}.bin"
dirpath = os.path.join(str(db_path), names[i])
os.makedirs(dirpath, exist_ok=True)
filepath = os.path.join(str(db_path), names[i], filename)
gt_points = points[point_indices[:, i]]
gt_points[:, :3] -= gt_boxes[i, :3]
with open(filepath, "w") as f:
try:
gt_points[:, :point_features].tofile(f)
except:
print("process {} files".format(index))
break
if (used_classes is None) or names[i] in used_classes:
if relative_path:
db_dump_path = os.path.join(db_path.stem, names[i],
filename)
else:
db_dump_path = str(filepath)
db_info = {
"name": names[i],
"path": db_dump_path,
"image_idx": image_idx,
"gt_idx": i,
"box3d_lidar": gt_boxes[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]
if names[i] in all_db_infos:
all_db_infos[names[i]].append(db_info)
else:
all_db_infos[names[i]] = [db_info]
print("dataset length: ", len(dataset))
for k, v in all_db_infos.items():
print(f"load {len(v)} {k} database infos")
with open(dbinfo_path, "wb") as f:
pickle.dump(all_db_infos, f)
def prep_pointcloud_rpn(
input_dict,
root_path,
task_class_names=[],
prep_cfg=None,
db_sampler=None,
remove_outside_points=False,
training=True,
num_point_features=4,
random_crop=False,
reference_detections=None,
out_dtype=np.float32,
min_points_in_gt=-1,
logger=None,
):
"""
convert point cloud to voxels, create targets if ground truths exists.
input_dict format: dataset.get_sensor_data format
"""
assert prep_cfg is not None
remove_environment = prep_cfg.REMOVE_UNKOWN_EXAMPLES
if training:
remove_unknown = prep_cfg.REMOVE_UNKOWN_EXAMPLES
gt_rotation_noise = prep_cfg.GT_ROT_NOISE
gt_loc_noise_std = prep_cfg.GT_LOC_NOISE
global_rotation_noise = prep_cfg.GLOBAL_ROT_NOISE
global_scaling_noise = prep_cfg.GLOBAL_SCALE_NOISE
global_random_rot_range = prep_cfg.GLOBAL_ROT_PER_OBJ_RANGE
global_translate_noise_std = prep_cfg.GLOBAL_TRANS_NOISE
gt_points_drop = prep_cfg.GT_DROP_PERCENTAGE
gt_drop_max_keep = prep_cfg.GT_DROP_MAX_KEEP_POINTS
remove_points_after_sample = prep_cfg.REMOVE_POINTS_AFTER_SAMPLE
class_names = list(itertools.chain(*task_class_names))
# points_only = input_dict["lidar"]["points"]
# times = input_dict["lidar"]["times"]
# points = np.hstack([points_only, times])
points = input_dict["lidar"]["points"]
if training:
anno_dict = input_dict["lidar"]["annotations"]
gt_dict = {
"gt_boxes": anno_dict["boxes"],
"gt_names": np.array(anno_dict["names"]).reshape(-1),
}
if "difficulty" not in anno_dict:
difficulty = np.zeros([anno_dict["boxes"].shape[0]],
dtype=np.int32)
gt_dict["difficulty"] = difficulty
else:
gt_dict["difficulty"] = anno_dict["difficulty"]
# if use_group_id and "group_ids" in anno_dict:
# group_ids = anno_dict["group_ids"]
# gt_dict["group_ids"] = group_ids
calib = None
if "calib" in input_dict:
calib = input_dict["calib"]
if reference_detections is not None:
assert calib is not None and "image" in input_dict
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.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:
assert calib is not None
image_shape = input_dict["image"]["image_shape"]
points = box_np_ops.remove_outside_points(points, calib["rect"],
calib["Trv2c"], calib["P2"],
image_shape)
if remove_environment is True and training:
selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
_dict_select(gt_dict, selected)
masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
points = points[masks.any(-1)]
if training:
selected = kitti.drop_arrays_by_name(gt_dict["gt_names"],
["DontCare", "ignore"])
_dict_select(gt_dict, selected)
if remove_unknown:
remove_mask = gt_dict["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)
_dict_select(gt_dict, keep_mask)
gt_dict.pop("difficulty")
gt_boxes_mask = np.array(
[n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
# db_sampler = None
if db_sampler is not None:
group_ids = None
# if "group_ids" in gt_dict:
# group_ids = gt_dict["group_ids"]
sampled_dict = db_sampler.sample_all(
root_path,
gt_dict["gt_boxes"],
gt_dict["gt_names"],
num_point_features,
random_crop,
gt_group_ids=group_ids,
calib=calib,
)
if sampled_dict is not None:
sampled_gt_names = sampled_dict["gt_names"]
sampled_gt_boxes = sampled_dict["gt_boxes"]
sampled_points = sampled_dict["points"]
sampled_gt_masks = sampled_dict["gt_masks"]
gt_dict["gt_names"] = np.concatenate(
[gt_dict["gt_names"], sampled_gt_names], axis=0)
gt_dict["gt_boxes"] = np.concatenate(
[gt_dict["gt_boxes"], sampled_gt_boxes])
gt_boxes_mask = np.concatenate(
[gt_boxes_mask, sampled_gt_masks], axis=0)
# if group_ids is not None:
# sampled_group_ids = sampled_dict["group_ids"]
# gt_dict["group_ids"] = np.concatenate(
# [gt_dict["group_ids"], sampled_group_ids])
if remove_points_after_sample:
masks = box_np_ops.points_in_rbbox(points,
sampled_gt_boxes)
points = points[np.logical_not(masks.any(-1))]
points = np.concatenate([sampled_points, points], axis=0)
# group_ids = None
# if "group_ids" in gt_dict:
# group_ids = gt_dict["group_ids"]
prep.noise_per_object_v3_(
gt_dict["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=None,
num_try=100,
)
# should remove unrelated objects after noise per object
# for k, v in gt_dict.items():
# print(k, v.shape)
_dict_select(gt_dict, gt_boxes_mask)
gt_classes = np.array(
[class_names.index(n) + 1 for n in gt_dict["gt_names"]],
dtype=np.int32)
gt_dict["gt_classes"] = gt_classes
gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
points)
gt_dict["gt_boxes"], points = prep.global_rotation(
gt_dict["gt_boxes"], points, rotation=global_rotation_noise)
gt_dict["gt_boxes"], points = prep.global_scaling_v2(
gt_dict["gt_boxes"], points, *global_scaling_noise)
prep.global_translate_(gt_dict["gt_boxes"], points,
global_translate_noise_std)
task_masks = []
flag = 0
for class_name in task_class_names:
task_masks.append([
np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 +
flag) for i in class_name
])
flag += len(class_name)
task_boxes = []
task_classes = []
task_names = []
flag2 = 0
for idx, mask in enumerate(task_masks):
task_box = []
task_class = []
task_name = []
for m in mask:
task_box.append(gt_dict["gt_boxes"][m])
task_class.append(gt_dict["gt_classes"][m] - flag2)
task_name.append(gt_dict["gt_names"][m])
task_boxes.append(np.concatenate(task_box, axis=0))
task_classes.append(np.concatenate(task_class))
task_names.append(np.concatenate(task_name))
flag2 += len(mask)
for task_box in task_boxes:
# limit rad to [-pi, pi]
task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1],
offset=0.5,
period=2 * np.pi)
# print(gt_dict.keys())
gt_dict["gt_classes"] = task_classes
gt_dict["gt_names"] = task_names
gt_dict["gt_boxes"] = task_boxes
example = {
"pts_input": points,
"pts_rect": None,
"pts_features": None,
"gt_boxes3d": gt_dict["gt_boxes"],
"rpn_cls_label": [],
"rpn_reg_label": [],
}
if calib is not None:
example["calib"] = calib
return example
def prep_sequence_pointcloud(
input_dict,
root_path,
voxel_generator,
target_assigners,
prep_cfg=None,
db_sampler=None,
remove_outside_points=False,
training=True,
create_targets=True,
num_point_features=4,
anchor_cache=None,
random_crop=False,
reference_detections=None,
out_size_factor=2,
out_dtype=np.float32,
min_points_in_gt=-1,
logger=None,
):
"""
convert point cloud to voxels, create targets if ground truths exists.
input_dict format: dataset.get_sensor_data format
"""
assert prep_cfg is not None
remove_environment = prep_cfg.REMOVE_ENVIRONMENT
max_voxels = prep_cfg.MAX_VOXELS_NUM
shuffle_points = prep_cfg.SHUFFLE
anchor_area_threshold = prep_cfg.ANCHOR_AREA_THRES
if training:
remove_unknown = prep_cfg.REMOVE_UNKOWN_EXAMPLES
gt_rotation_noise = prep_cfg.GT_ROT_NOISE
gt_loc_noise_std = prep_cfg.GT_LOC_NOISE
global_rotation_noise = prep_cfg.GLOBAL_ROT_NOISE
global_scaling_noise = prep_cfg.GLOBAL_SCALE_NOISE
global_random_rot_range = prep_cfg.GLOBAL_ROT_PER_OBJ_RANGE
global_translate_noise_std = prep_cfg.GLOBAL_TRANS_NOISE
gt_points_drop = prep_cfg.GT_DROP_PERCENTAGE
gt_drop_max_keep = prep_cfg.GT_DROP_MAX_KEEP_POINTS
remove_points_after_sample = prep_cfg.REMOVE_POINTS_AFTER_SAMPLE
min_points_in_gt = prep_cfg.get("MIN_POINTS_IN_GT", -1)
task_class_names = [
target_assigner.classes for target_assigner in target_assigners
]
class_names = list(itertools.chain(*task_class_names))
# points_only = input_dict["lidar"]["points"]
# times = input_dict["lidar"]["times"]
# points = np.hstack([points_only, times])
try:
points = input_dict["current_frame"]["lidar"]["combined"]
except Exception:
points = input_dict["current_frame"]["lidar"]["points"]
keyframe_points = input_dict["keyframe"]["lidar"]["combined"]
if training:
anno_dict = input_dict["current_frame"]["lidar"]["annotations"]
gt_dict = {
"gt_boxes": anno_dict["boxes"],
"gt_names": np.array(anno_dict["names"]).reshape(-1),
}
if "difficulty" not in anno_dict:
difficulty = np.zeros([anno_dict["boxes"].shape[0]],
dtype=np.int32)
gt_dict["difficulty"] = difficulty
else:
gt_dict["difficulty"] = anno_dict["difficulty"]
# if use_group_id and "group_ids" in anno_dict:
# group_ids = anno_dict["group_ids"]
# gt_dict["group_ids"] = group_ids
calib = None
if "calib" in input_dict:
calib = input_dict["current_frame"]["calib"]
if reference_detections is not None:
assert calib is not None and "image" in input_dict["current_frame"]
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.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:
assert calib is not None
image_shape = input_dict["current_frame"]["image"]["image_shape"]
points = box_np_ops.remove_outside_points(points, calib["rect"],
calib["Trv2c"], calib["P2"],
image_shape)
if remove_environment is True and training:
selected = kitti.keep_arrays_by_name(gt_names, target_assigner.classes)
_dict_select(gt_dict, selected)
masks = box_np_ops.points_in_rbbox(points, gt_dict["gt_boxes"])
points = points[masks.any(-1)]
if training:
# boxes_lidar = gt_dict["gt_boxes"]
# cv2.imshow('pre-noise', bev_map)
selected = kitti.drop_arrays_by_name(gt_dict["gt_names"],
["DontCare", "ignore"])
_dict_select(gt_dict, selected)
if remove_unknown:
remove_mask = gt_dict["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)
_dict_select(gt_dict, keep_mask)
gt_dict.pop("difficulty")
if min_points_in_gt > 0:
# points_count_rbbox takes 10ms with 10 sweeps nuscenes data
point_counts = box_np_ops.points_count_rbbox(
points, gt_dict["gt_boxes"])
mask = point_counts >= min_points_in_gt
_dict_select(gt_dict, mask)
gt_boxes_mask = np.array(
[n in class_names for n in gt_dict["gt_names"]], dtype=np.bool_)
# db_sampler = None
if db_sampler is not None:
group_ids = None
# if "group_ids" in gt_dict:
# group_ids = gt_dict["group_ids"]
sampled_dict = db_sampler.sample_all(
root_path,
gt_dict["gt_boxes"],
gt_dict["gt_names"],
num_point_features,
random_crop,
gt_group_ids=group_ids,
calib=calib,
)
if sampled_dict is not None:
sampled_gt_names = sampled_dict["gt_names"]
sampled_gt_boxes = sampled_dict["gt_boxes"]
sampled_points = sampled_dict["points"]
sampled_gt_masks = sampled_dict["gt_masks"]
gt_dict["gt_names"] = np.concatenate(
[gt_dict["gt_names"], sampled_gt_names], axis=0)
gt_dict["gt_boxes"] = np.concatenate(
[gt_dict["gt_boxes"], sampled_gt_boxes])
gt_boxes_mask = np.concatenate(
[gt_boxes_mask, sampled_gt_masks], axis=0)
# if group_ids is not None:
# sampled_group_ids = sampled_dict["group_ids"]
# gt_dict["group_ids"] = np.concatenate(
# [gt_dict["group_ids"], sampled_group_ids])
if remove_points_after_sample:
masks = box_np_ops.points_in_rbbox(points,
sampled_gt_boxes)
points = points[np.logical_not(masks.any(-1))]
points = np.concatenate([sampled_points, points], axis=0)
pc_range = voxel_generator.point_cloud_range
# group_ids = None
# if "group_ids" in gt_dict:
# group_ids = gt_dict["group_ids"]
# prep.noise_per_object_v3_(
# gt_dict["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)
# should remove unrelated objects after noise per object
# for k, v in gt_dict.items():
# print(k, v.shape)
_dict_select(gt_dict, gt_boxes_mask)
gt_classes = np.array(
[class_names.index(n) + 1 for n in gt_dict["gt_names"]],
dtype=np.int32)
gt_dict["gt_classes"] = gt_classes
# concatenate
points_current = points.shape[0]
points_keyframe = keyframe_points.shape[0]
points = np.concatenate((points, keyframe_points), axis=0)
# data aug
gt_dict["gt_boxes"], points = prep.random_flip(gt_dict["gt_boxes"],
points)
gt_dict["gt_boxes"], points = prep.global_rotation(
gt_dict["gt_boxes"], points, rotation=global_rotation_noise)
gt_dict["gt_boxes"], points = prep.global_scaling_v2(
gt_dict["gt_boxes"], points, *global_scaling_noise)
prep.global_translate_(gt_dict["gt_boxes"], points,
global_translate_noise_std)
# slice
points_keyframe = points[points_current:, :]
points = points[:points_current, :]
bv_range = voxel_generator.point_cloud_range[[0, 1, 3, 4]]
mask = prep.filter_gt_box_outside_range(gt_dict["gt_boxes"], bv_range)
_dict_select(gt_dict, mask)
task_masks = []
flag = 0
for class_name in task_class_names:
task_masks.append([
np.where(gt_dict["gt_classes"] == class_name.index(i) + 1 +
flag) for i in class_name
])
flag += len(class_name)
task_boxes = []
task_classes = []
task_names = []
flag2 = 0
for idx, mask in enumerate(task_masks):
task_box = []
task_class = []
task_name = []
for m in mask:
task_box.append(gt_dict["gt_boxes"][m])
task_class.append(gt_dict["gt_classes"][m] - flag2)
task_name.append(gt_dict["gt_names"][m])
task_boxes.append(np.concatenate(task_box, axis=0))
task_classes.append(np.concatenate(task_class))
task_names.append(np.concatenate(task_name))
flag2 += len(mask)
for task_box in task_boxes:
# limit rad to [-pi, pi]
task_box[:, -1] = box_np_ops.limit_period(task_box[:, -1],
offset=0.5,
period=2 * np.pi)
# print(gt_dict.keys())
gt_dict["gt_classes"] = task_classes
gt_dict["gt_names"] = task_names
gt_dict["gt_boxes"] = task_boxes
# if shuffle_points:
# # shuffle is a little slow.
# np.random.shuffle(points)
# [0, -40, -3, 70.4, 40, 1]
voxel_size = voxel_generator.voxel_size
pc_range = voxel_generator.point_cloud_range
grid_size = voxel_generator.grid_size
# [352, 400]
# points = points[:int(points.shape[0] * 0.1), :]
voxels, coordinates, num_points = voxel_generator.generate(
points, max_voxels)
# res = voxel_generator.generate(
# points, max_voxels)
# voxels = res["voxels"]
# coordinates = res["coordinates"]
# num_points = res["num_points_per_voxel"]
num_voxels = np.array([voxels.shape[0]], dtype=np.int64)
# key frame voxel
keyframe_info = voxel_generator.generate(keyframe_points, max_voxels)
keyframe_info = keyframe_voxels, keyframe_coordinates, keyframe_num_points
keyframe_num_voxels = np.array([keyframe_voxels.shape[0]], dtype=np.int64)
example = {
"voxels": voxels,
"num_points": num_points,
"points": points,
"coordinates": coordinates,
"num_voxels": num_voxels,
}
example_keyframe = {
"voxels": keyframe_voxels,
"num_points": keyframe_num_points,
"points": keyframe_points,
"coordinates": keyframe_coordinates,
"num_voxels": keyframe_num_voxels,
}
if training:
example["gt_boxes"] = gt_dict["gt_boxes"]
if calib is not None:
example["calib"] = calib
feature_map_size = grid_size[:2] // out_size_factor
feature_map_size = [*feature_map_size, 1][::-1]
if anchor_cache is not None:
anchorss = anchor_cache["anchors"]
anchors_bvs = anchor_cache["anchors_bv"]
anchors_dicts = anchor_cache["anchors_dict"]
else:
rets = [
target_assigner.generate_anchors(feature_map_size)
for target_assigner in target_assigners
]
anchorss = [ret["anchors"].reshape([-1, 7]) for ret in rets]
anchors_dicts = [
target_assigner.generate_anchors_dict(feature_map_size)
for target_assigner in target_assigners
]
anchors_bvs = [
box_np_ops.rbbox2d_to_near_bbox(anchors[:, [0, 1, 3, 4, 6]])
for anchors in anchorss
]
example["anchors"] = anchorss
if anchor_area_threshold >= 0:
example["anchors_mask"] = []
for idx, anchors_bv in enumerate(anchors_bvs):
anchors_mask = None
# slow with high resolution. recommend disable this forever.
coors = coordinates
dense_voxel_map = box_np_ops.sparse_sum_for_anchors_mask(
coors, tuple(grid_size[::-1][1:]))
dense_voxel_map = dense_voxel_map.cumsum(0)
dense_voxel_map = dense_voxel_map.cumsum(1)
anchors_area = box_np_ops.fused_get_anchors_area(
dense_voxel_map, anchors_bv, voxel_size, pc_range, grid_size)
anchors_mask = anchors_area > anchor_area_threshold
# example['anchors_mask'] = anchors_mask.astype(np.uint8)
example["anchors_mask"].append(anchors_mask)
example_sequences = {}
example_sequences["current_frame"] = example
example_sequences["keyframe"] = example_keyframe
if not training:
return example_sequences
# voxel_labels = box_np_ops.assign_label_to_voxel(gt_boxes, coordinates,
# voxel_size, coors_range)
"""
example.update({
'gt_boxes': gt_boxes.astype(out_dtype),
'num_gt': np.array([gt_boxes.shape[0]]),
# 'voxel_labels': voxel_labels,
})
"""
if create_targets:
targets_dicts = []
for idx, target_assigner in enumerate(target_assigners):
if "anchors_mask" in example:
anchors_mask = example["anchors_mask"][idx]
else:
anchors_mask = None
targets_dict = target_assigner.assign_v2(
anchors_dicts[idx],
gt_dict["gt_boxes"][idx],
anchors_mask,
gt_classes=gt_dict["gt_classes"][idx],
gt_names=gt_dict["gt_names"][idx],
)
targets_dicts.append(targets_dict)
example_sequences["current_frame"].update({
"labels":
[targets_dict["labels"] for targets_dict in targets_dicts],
"reg_targets":
[targets_dict["bbox_targets"] for targets_dict in targets_dicts],
"reg_weights": [
targets_dict["bbox_outside_weights"]
for targets_dict in targets_dicts
],
})
return example_sequences