def sample_all(
self,
root_path,
gt_boxes,
gt_names,
num_point_features,
random_crop=False,
gt_group_ids=None,
calib=None,
targeted_class_names=None,
with_road_plane_cam=None,
):
'''
This func aims to sample some gt-boxes and corresponding points to perform gt augmentation;
notice that the non-targeted gt boxes (like pedestrian) have been considered into collision test;
notice that the points in corresponding gt-box are read from pre-saved gt database;
'''
# record the num of gt-aug samples with a dict and a list
sampled_num_dict = {}
sample_num_per_class = []
for class_name, max_sample_num in zip(
self._sample_classes,
self._sample_max_nums): # actual only once for ['Car': 15]
sampled_num = int(max_sample_num -
np.sum([n == class_name for n in gt_names]))
#sampled_num = int(max_sample_num - np.sum([name in targeted_class_names for name in gt_names]))
sampled_num = np.round(self._rate * sampled_num).astype(np.int64)
sampled_num_dict[class_name] = sampled_num
sample_num_per_class.append(sampled_num)
sampled = []
sampled_gt_boxes = []
all_gt_boxes = gt_boxes
# gt-augmentation: sample gt boxes and add them to current gt_boxes.
# todo: we may sample box one by one to ensure num of gt-boxes is fulfilled.
for class_name, sampled_num in zip(self._sample_classes,
sample_num_per_class):
# if sampled_num > 0:
# sampled_objects = self.sample_class_v2(class_name, sampled_num, all_gt_boxes)
# sampled += sampled_objects
# if len(sampled_objects) > 0:
# if len(sampled_objects) == 1:
# sampled_boxes = sampled_objects[0]["box3d_lidar"][np.newaxis, ...]
# else:
# sampled_boxes = np.stack([s["box3d_lidar"] for s in sampled_objects], axis=0)
# sampled_gt_boxes += [sampled_boxes]
# all_gt_boxes = np.concatenate([all_gt_boxes, sampled_boxes], axis=0)
# ensure final num_boxes fulfill the num requirement after collision test.
times = 0
while sampled_num > 0 and times < 2:
sampled_objects = self.sample_class_v2(class_name, sampled_num,
all_gt_boxes)
sampled += sampled_objects
if len(sampled_objects) > 0:
if len(sampled_objects) == 1:
sampled_boxes = sampled_objects[0]["box3d_lidar"][
np.newaxis, ...]
else:
sampled_boxes = np.stack(
[s["box3d_lidar"] for s in sampled_objects],
axis=0)
sampled_gt_boxes += [sampled_boxes]
all_gt_boxes = np.concatenate(
[all_gt_boxes, sampled_boxes], axis=0)
sampled_num -= len(sampled_objects)
times += 1
if len(sampled) > 0:
''' get points in sampled gt_boxes '''
sampled_gt_boxes = np.concatenate(sampled_gt_boxes, axis=0)
num_sampled = len(sampled)
s_points_list = []
# get points in sampled gt-boxes from pre-generated gt database.
for info in sampled:
try:
s_points = np.fromfile(
str(pathlib.Path(root_path) / info["path"]),
dtype=np.float32).reshape(-1, num_point_features)
# gt_points are saved with relative distance; so need to recover by adding box center.
s_points[:, :3] += info["box3d_lidar"][:3]
if with_road_plane_cam is not None:
a, b, c, d = with_road_plane_cam
box3d_cam_center = info['box3d_cam'][
0:3] # x,y,z with cam coord. bottom center.
cur_height_cam = (-d - a * box3d_cam_center[0] -
c * box3d_cam_center[2]) / b
move_height_cam = info['box3d_cam'][
1] - cur_height_cam # cal y dist, > 0: move up, < 0: move down.
s_points[:, 2] += move_height_cam
index = sampled.index(info)
sampled_gt_boxes[index, 2] += move_height_cam
# random drop points in gt_boxes to make model more robust.
if self.gt_point_random_drop > 0:
num_point = s_points.shape[0]
if num_point > 10:
drop_num = int(
np.random.uniform(0, self.gt_point_random_drop)
* num_point)
choice = np.random.choice(np.arange(num_point),
num_point - drop_num,
replace=False)
s_points = s_points[choice]
s_points_list.append(s_points)
except Exception:
print(info["path"])
continue
'''todo: do something about random crop'''
if random_crop: # False
s_points_list_new = []
assert calib is not None
rect = calib["rect"]
Trv2c = calib["Trv2c"]
P2 = calib["P2"]
gt_bboxes = box_np_ops.box3d_to_bbox(sampled_gt_boxes, rect,
Trv2c, P2)
crop_frustums = prep.random_crop_frustum(
gt_bboxes, rect, Trv2c, P2)
for i in range(crop_frustums.shape[0]):
s_points = s_points_list[i]
mask = prep.mask_points_in_corners(
s_points, crop_frustums[i:i + 1]).reshape(-1)
num_remove = np.sum(mask)
if num_remove > 0 and (s_points.shape[0] -
num_remove) > 15:
s_points = s_points[np.logical_not(mask)]
s_points_list_new.append(s_points)
s_points_list = s_points_list_new
ret = {
"gt_names":
np.array([s["name"] for s in sampled]),
"difficulty":
np.array([s["difficulty"] for s in sampled]),
"gt_boxes":
sampled_gt_boxes,
"points":
np.concatenate(s_points_list, axis=0),
"gt_masks":
np.ones((num_sampled, ), dtype=np.bool_),
"group_ids":
np.arange(gt_boxes.shape[0], gt_boxes.shape[0] + len(sampled)),
}
else:
ret = None
return ret
def sample_all(
self,
root_path,
gt_boxes,
gt_names,
num_point_features,
random_crop=False,
gt_group_ids=None,
calib=None,
):
# record the num of gt-aug samples with a dict and a list
sampled_num_dict = {}
sample_num_per_class = []
for class_name, max_sample_num in zip(
self._sample_classes,
self._sample_max_nums): # actual only once for ['Car': 15]
sampled_num = int(max_sample_num -
np.sum([n == class_name for n in gt_names]))
sampled_num = np.round(self._rate * sampled_num).astype(np.int64)
sampled_num_dict[class_name] = sampled_num
sample_num_per_class.append(sampled_num)
sampled_groups = self._sample_classes
if self._use_group_sampling: # False
assert gt_group_ids is not None
sampled_groups = []
sample_num_per_class = []
for group_name, class_names in self._group_name_to_names:
sampled_nums_group = [sampled_num_dict[n] for n in class_names]
sampled_num = np.max(sampled_nums_group)
sample_num_per_class.append(sampled_num)
sampled_groups.append(group_name)
total_group_ids = gt_group_ids
sampled = []
sampled_gt_boxes = []
avoid_coll_boxes = gt_boxes
# gt-augmentation: sample gt boxes and add them to current gt_boxes.
# todo: we may sample box one by one to ensure num of gt-boxes is fulfilled.
for class_name, sampled_num in zip(sampled_groups,
sample_num_per_class):
if sampled_num > 0:
if self._use_group_sampling: # False
sampled_cls = self.sample_group(class_name, sampled_num,
avoid_coll_boxes,
total_group_ids)
else:
sampled_cls = self.sample_class_v2(class_name, sampled_num,
avoid_coll_boxes)
sampled += sampled_cls
if len(sampled_cls) > 0:
if len(sampled_cls) == 1: # True
sampled_gt_box = sampled_cls[0]["box3d_lidar"][
np.newaxis, ...]
else:
sampled_gt_box = np.stack(
[s["box3d_lidar"] for s in sampled_cls], axis=0)
sampled_gt_boxes += [sampled_gt_box]
avoid_coll_boxes = np.concatenate(
[avoid_coll_boxes, sampled_gt_box], axis=0)
if self._use_group_sampling: # False
if len(sampled_cls) == 1:
sampled_group_ids = np.array(
sampled_cls[0]["group_id"])[np.newaxis, ...]
else:
sampled_group_ids = np.stack(
[s["group_id"] for s in sampled_cls], axis=0)
total_group_ids = np.concatenate(
[total_group_ids, sampled_group_ids], axis=0)
if len(sampled) > 0:
''' get points in sampled gt_boxes '''
sampled_gt_boxes = np.concatenate(sampled_gt_boxes, axis=0)
num_sampled = len(sampled)
s_points_list = []
# get points in sampled gt-boxes.
for info in sampled: # info.keys: ['name', 'path', 'image_idx', 'gt_idx', 'box3d_lidar', 'num_points_in_gt', 'difficulty', 'group_id']
try:
# TODO fix point read error
s_points = np.fromfile(
str(pathlib.Path(root_path) / info["path"]),
dtype=np.float32).reshape(-1, num_point_features)
if "rot_transform" in info: # False
rot = info["rot_transform"]
s_points[:, :
3] = box_np_ops.rotation_points_single_angle(
s_points[:, :4], rot, axis=2)
# gt_points are saved with relative distance; so need to recover by adding box center.
s_points[:, :3] += info["box3d_lidar"][:3]
s_points_list.append(s_points)
except Exception:
print(info["path"])
continue
'''todo: do something about random crop'''
if random_crop: # False
s_points_list_new = []
assert calib is not None
rect = calib["rect"]
Trv2c = calib["Trv2c"]
P2 = calib["P2"]
gt_bboxes = box_np_ops.box3d_to_bbox(sampled_gt_boxes, rect,
Trv2c, P2)
crop_frustums = prep.random_crop_frustum(
gt_bboxes, rect, Trv2c, P2)
for i in range(crop_frustums.shape[0]):
s_points = s_points_list[i]
mask = prep.mask_points_in_corners(
s_points, crop_frustums[i:i + 1]).reshape(-1)
num_remove = np.sum(mask)
if num_remove > 0 and (s_points.shape[0] -
num_remove) > 15:
s_points = s_points[np.logical_not(mask)]
s_points_list_new.append(s_points)
s_points_list = s_points_list_new
ret = {
"gt_names": np.array([s["name"] for s in sampled]),
"difficulty": np.array([s["difficulty"] for s in sampled]),
"gt_boxes": sampled_gt_boxes,
"points": np.concatenate(s_points_list, axis=0),
"gt_masks": np.ones((num_sampled, ), dtype=np.bool_),
}
if self._use_group_sampling:
ret["group_ids"] = np.array([s["group_id"] for s in sampled])
else:
ret["group_ids"] = np.arange(gt_boxes.shape[0],
gt_boxes.shape[0] + len(sampled))
else:
ret = None
return ret
def sample_all(
self,
root_path,
gt_boxes,
gt_names,
num_point_features,
random_crop=False,
gt_group_ids=None,
calib=None,
road_planes=None,
):
sampled_num_dict = {}
sample_num_per_class = []
for class_name, max_sample_num in zip(self._sample_classes,
self._sample_max_nums):
sampled_num = int(max_sample_num -
np.sum([n == class_name for n in gt_names]))
sampled_num = np.round(self._rate * sampled_num).astype(np.int64)
sampled_num_dict[class_name] = sampled_num
sample_num_per_class.append(sampled_num)
sampled_groups = self._sample_classes
if self._use_group_sampling:
assert gt_group_ids is not None
sampled_groups = []
sample_num_per_class = []
for group_name, class_names in self._group_name_to_names:
sampled_nums_group = [sampled_num_dict[n] for n in class_names]
sampled_num = np.max(sampled_nums_group)
sample_num_per_class.append(sampled_num)
sampled_groups.append(group_name)
total_group_ids = gt_group_ids
sampled = []
sampled_gt_boxes = []
avoid_coll_boxes = gt_boxes
for class_name, sampled_num in zip(sampled_groups,
sample_num_per_class):
if sampled_num > 0:
if self._use_group_sampling:
sampled_cls = self.sample_group(class_name, sampled_num,
avoid_coll_boxes,
total_group_ids)
else:
sampled_cls = self.sample_class_v2(class_name, sampled_num,
avoid_coll_boxes)
sampled += sampled_cls
if len(sampled_cls) > 0:
if len(sampled_cls) == 1:
sampled_gt_box = sampled_cls[0]["box3d_lidar"][
np.newaxis, ...]
else:
sampled_gt_box = np.stack(
[s["box3d_lidar"] for s in sampled_cls], axis=0)
sampled_gt_boxes += [sampled_gt_box]
avoid_coll_boxes = np.concatenate(
[avoid_coll_boxes, sampled_gt_box], axis=0)
if self._use_group_sampling:
if len(sampled_cls) == 1:
sampled_group_ids = np.array(
sampled_cls[0]["group_id"])[np.newaxis, ...]
else:
sampled_group_ids = np.stack(
[s["group_id"] for s in sampled_cls], axis=0)
total_group_ids = np.concatenate(
[total_group_ids, sampled_group_ids], axis=0)
if len(sampled) > 0:
sampled_gt_boxes = np.concatenate(sampled_gt_boxes, axis=0)
num_sampled = len(sampled)
s_points_list = []
for info in sampled:
try:
s_points = np.fromfile(
str(pathlib.Path(root_path) / info["path"]),
dtype=np.float32).reshape(-1, num_point_features)
if "rot_transform" in info:
rot = info["rot_transform"]
s_points[:, :
3] = box_np_ops.rotation_points_single_angle(
s_points[:, :4], rot, axis=2)
s_points[:, :3] += info["box3d_lidar"][:3]
s_points_list.append(s_points)
# print(pathlib.Path(info["path"]).stem)
except Exception:
print(str(pathlib.Path(root_path) / info["path"]))
continue
if random_crop:
s_points_list_new = []
assert calib is not None
rect = calib["rect"]
Trv2c = calib["Trv2c"]
P2 = calib["P2"]
gt_bboxes = box_np_ops.box3d_to_bbox(sampled_gt_boxes, rect,
Trv2c, P2)
crop_frustums = prep.random_crop_frustum(
gt_bboxes, rect, Trv2c, P2)
for i in range(crop_frustums.shape[0]):
s_points = s_points_list[i]
mask = prep.mask_points_in_corners(
s_points, crop_frustums[i:i + 1]).reshape(-1)
num_remove = np.sum(mask)
if num_remove > 0 and (s_points.shape[0] -
num_remove) > 15:
s_points = s_points[np.logical_not(mask)]
s_points_list_new.append(s_points)
s_points_list = s_points_list_new
ret = {
"gt_names": np.array([s["name"] for s in sampled]),
"difficulty": np.array([s["difficulty"] for s in sampled]),
"gt_boxes": sampled_gt_boxes,
"points": np.concatenate(s_points_list, axis=0),
"gt_masks": np.ones((num_sampled, ), dtype=np.bool_),
}
if self._use_group_sampling:
ret["group_ids"] = np.array([s["group_id"] for s in sampled])
else:
ret["group_ids"] = np.arange(gt_boxes.shape[0],
gt_boxes.shape[0] + len(sampled))
else:
ret = None
return ret
def sample_all(
self,
root_path,
gt_boxes,
gt_names,
num_point_features,
random_crop=False,
gt_group_ids=None,
calib=None,
road_planes=None,
):
sampled_num_dict = {}
sample_num_per_class = []
for class_name, max_sample_num in zip(self._sample_classes,
self._sample_max_nums):
sampled_num = int(max_sample_num -
np.sum([n == class_name for n in gt_names]))
sampled_num = np.round(self._rate * sampled_num).astype(np.int64)
sampled_num_dict[class_name] = sampled_num
sample_num_per_class.append(sampled_num)
sampled_groups = self._sample_classes
if self._use_group_sampling:
assert gt_group_ids is not None
sampled_groups = []
sample_num_per_class = []
for group_name, class_names in self._group_name_to_names:
sampled_nums_group = [sampled_num_dict[n] for n in class_names]
sampled_num = np.max(sampled_nums_group)
sample_num_per_class.append(sampled_num)
sampled_groups.append(group_name)
total_group_ids = gt_group_ids
sampled = []
sampled_gt_boxes = []
avoid_coll_boxes = gt_boxes
for class_name, sampled_num in zip(sampled_groups,
sample_num_per_class):
if sampled_num > 0:
if self._use_group_sampling:
sampled_cls = self.sample_group(class_name, sampled_num,
avoid_coll_boxes,
total_group_ids)
else:
sampled_cls = self.sample_class_v2(class_name, sampled_num,
avoid_coll_boxes)
sampled += sampled_cls
if len(sampled_cls) > 0:
if len(sampled_cls) == 1:
sampled_gt_box = sampled_cls[0]["box3d_lidar"][
np.newaxis, ...]
else:
sampled_gt_box = np.stack(
[s["box3d_lidar"] for s in sampled_cls], axis=0)
sampled_gt_boxes += [sampled_gt_box]
avoid_coll_boxes = np.concatenate(
[avoid_coll_boxes, sampled_gt_box], axis=0)
if self._use_group_sampling:
if len(sampled_cls) == 1:
sampled_group_ids = np.array(
sampled_cls[0]["group_id"])[np.newaxis, ...]
else:
sampled_group_ids = np.stack(
[s["group_id"] for s in sampled_cls], axis=0)
total_group_ids = np.concatenate(
[total_group_ids, sampled_group_ids], axis=0)
if len(sampled) > 0:
sampled_gt_boxes = np.concatenate(sampled_gt_boxes, axis=0)
if road_planes is not None:
# Only support KITTI
# image plane
assert False, "Not correct yet!"
a, b, c, d = road_planes
center = sampled_gt_boxes[:, :3]
center[:, 2] -= sampled_gt_boxes[:, 5] / 2
center_cam = box_np_ops.lidar_to_camera(
center, calib["rect"], calib["Trv2c"])
cur_height_cam = (-d - a * center_cam[:, 0] -
c * center_cam[:, 2]) / b
center_cam[:, 1] = cur_height_cam
lidar_tmp_point = box_np_ops.camera_to_lidar(
center_cam, calib["rect"], calib["Trv2c"])
cur_lidar_height = lidar_tmp_point[:, 2]
# botom to middle center
# kitti [0.5, 0.5, 0] center to [0.5, 0.5, 0.5]
sampled_gt_boxes[:,
2] = cur_lidar_height + sampled_gt_boxes[:,
5] / 2
# mv_height = sampled_gt_boxes[:, 2] - cur_lidar_height
# sampled_gt_boxes[:, 2] -= mv_height
num_sampled = len(sampled)
s_points_list = []
for info in sampled:
try:
# TODO fix point read error
s_points = np.fromfile(
str(pathlib.Path(root_path) / info["path"]),
dtype=np.float32).reshape(-1, num_point_features)
# if not add_rgb_to_points:
# s_points = s_points[:, :4]
if "rot_transform" in info:
rot = info["rot_transform"]
s_points[:, :
3] = box_np_ops.rotation_points_single_angle(
s_points[:, :4], rot, axis=2)
s_points[:, :3] += info["box3d_lidar"][:3]
s_points_list.append(s_points)
# print(pathlib.Path(info["path"]).stem)
except Exception:
print(str(pathlib.Path(root_path) / info["path"]))
continue
# gt_bboxes = np.stack([s["bbox"] for s in sampled], axis=0)
# if np.random.choice([False, True], replace=False, p=[0.3, 0.7]):
# do random crop.
if random_crop:
s_points_list_new = []
assert calib is not None
rect = calib["rect"]
Trv2c = calib["Trv2c"]
P2 = calib["P2"]
gt_bboxes = box_np_ops.box3d_to_bbox(sampled_gt_boxes, rect,
Trv2c, P2)
crop_frustums = prep.random_crop_frustum(
gt_bboxes, rect, Trv2c, P2)
for i in range(crop_frustums.shape[0]):
s_points = s_points_list[i]
mask = prep.mask_points_in_corners(
s_points, crop_frustums[i:i + 1]).reshape(-1)
num_remove = np.sum(mask)
if num_remove > 0 and (s_points.shape[0] -
num_remove) > 15:
s_points = s_points[np.logical_not(mask)]
s_points_list_new.append(s_points)
s_points_list = s_points_list_new
ret = {
"gt_names": np.array([s["name"] for s in sampled]),
"difficulty": np.array([s["difficulty"] for s in sampled]),
"gt_boxes": sampled_gt_boxes,
"points": np.concatenate(s_points_list, axis=0),
"gt_masks": np.ones((num_sampled, ), dtype=np.bool_),
}
if self._use_group_sampling:
ret["group_ids"] = np.array([s["group_id"] for s in sampled])
else:
ret["group_ids"] = np.arange(gt_boxes.shape[0],
gt_boxes.shape[0] + len(sampled))
else:
ret = None
return ret