def flip_pc(self, label: WaymoLabel, pc: np.array,
calib: WaymoCalib) -> (WaymoLabel, np.array):
'''
flip point cloud along the y axis of the Kitti Lidar frame
inputs:
label: ground truth
pc: point cloud
calib:
'''
assert istype(label, "WaymoLabel")
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = WaymoLabel()
label_.data = []
for obj in label.data:
label_.data.append(WaymoObj(str(obj)))
# flip point cloud
if isinstance(pc_, dict):
for velo in pc_.keys():
pc_[velo][:, 1] *= -1
else:
pc_[:, 1] *= -1
# modify gt
for obj in label_.data:
obj.y *= -1
obj.ry *= -1
res_label = WaymoLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
def filt_label_by_num_of_pts(pc, calib, label, min_num_pts):
'''
@pc: np.ndarray
'''
if min_num_pts < 0:
return label
is_iter = isinstance(calib, list)
if istype(label, "CarlaLabel"):
res = CarlaLabel()
elif istype(label, "WaymoLabel"):
res = WaymoLabel()
elif istype(label, "KittiLabel"):
res = KittiLabel()
else:
raise NotImplementedError
res.data = [] # It is really important,
# otherwise the "for obj in label.data" of the next stage will raise a problem.
assert not isinstance(pc, dict)
if not is_iter:
for obj in label.data:
num_pts = obj.get_pts_idx(pc, calib).sum()
if num_pts > min_num_pts:
res.add_obj(obj)
else:
for obj, calib_ in zip(label.data, calib):
num_pts = obj.get_pts_idx(pc, calib_).sum()
if num_pts > min_num_pts:
res.add_obj(obj)
res.current_frame = label.current_frame
return res
def filt_label_by_range(label, valid_range, calib=None):
'''
@label: CarlaLabel
@valid_range: [min_x, min_y, min_z, max_x, max_y, max_z] FIMU
'''
min_x, min_y, min_z, max_x, max_y, max_z = valid_range
is_iter = isinstance(calib, list)
if istype(label, "CarlaLabel"):
res = CarlaLabel()
elif istype(label, "WaymoLabel"):
res = WaymoLabel()
elif istype(label, "KittiLabel"):
res = KittiLabel()
else:
raise NotImplementedError
res.data = [] # It is really important,
# otherwise the "for obj in label.data" of the next stage will raise a problem.
if not is_iter:
for obj in label.data:
if istype(label, "CarlaLabel") or istype(label, "WaymoLabel"):
imu_pt = np.array([obj.x, obj.y,
obj.z + obj.h / 2.0]).reshape(1, 3)
if (min_x <= imu_pt[0, 0] <= max_x
and min_y <= imu_pt[0, 1] <= max_y
and min_z <= imu_pt[0, 2] <= max_z):
res.add_obj(obj)
elif istype(label, "KittiLabel"):
bcenter_Fcam = np.array([obj.x, obj.y, obj.z]).reshape(1, 3)
bcenter_Flidar = calib.leftcam2lidar(bcenter_Fcam)
center_Flidar = bcenter_Flidar + np.array([0, 0, obj.h / 2.0
]).reshape(1, 3)
if (min_x <= center_Flidar[0, 0] <= max_x
and min_y <= center_Flidar[0, 1] <= max_y
and min_z <= center_Flidar[0, 2] <= max_z):
res.add_obj(obj)
else:
raise NotImplementedError
else:
for obj, calib_ in zip(label.data, calib):
if istype(label, "CarlaLabel") or istype(label, "WaymoLabel"):
imu_pt = np.array([obj.x, obj.y,
obj.z + obj.h / 2.0]).reshape(1, 3)
if (min_x <= imu_pt[0, 0] <= max_x
and min_y <= imu_pt[0, 1] <= max_y
and min_z <= imu_pt[0, 2] <= max_z):
res.add_obj(obj)
elif istype(label, "KittiLabel"):
bcenter_Fcam = np.array([obj.x, obj.y, obj.z]).reshape(1, 3)
bcenter_Flidar = calib_.leftcam2lidar(bcenter_Fcam)
center_Flidar = bcenter_Flidar + np.array([0, 0, obj.h / 2.0
]).reshape(1, 3)
if (min_x <= center_Flidar[0, 0] <= max_x
and min_y <= center_Flidar[0, 1] <= max_y
and min_z <= center_Flidar[0, 2] <= max_z):
res.add_obj(obj)
else:
raise NotImplementedError
res.current_frame = label.current_frame
return res
def keep(self, label: WaymoLabel, pc: np.array,
calib: WaymoCalib) -> (WaymoLabel, np.array):
assert istype(label, "WaymoLabel")
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = WaymoLabel()
label_.data = []
for obj in label.data:
label_.data.append(WaymoObj(str(obj)))
res_label = WaymoLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
def filt_label_by_cls(label, keep_classes):
if istype(label, "CarlaLabel"):
res = CarlaLabel()
elif istype(label, "WaymoLabel"):
res = WaymoLabel()
elif istype(label, "KittiLabel"):
res = KittiLabel()
else:
raise NotImplementedError
res.data = [] # It is really important,
# otherwise the "for obj in label.data" of the next stage will raise a problem.
for obj in label.data:
if obj.type in keep_classes:
res.add_obj(obj)
res.current_frame = label.current_frame
return res
def get_classes(label_dir: str, idx_list: list, dataset: str):
'''
Get all the classes by going through all the label files.
'''
res = []
for idx in idx_list:
if dataset == "carla":
label = CarlaLabel(
label_path=f"{label_dir}/{idx}.txt").read_label_file()
elif dataset == "waymo":
label = WaymoLabel(
label_path=f"{label_dir}/{idx}.txt").read_label_file()
elif dataset == "kitti":
label = KittiLabel(
label_path=f"{label_dir}/{idx}.txt").read_label_file()
else:
raise NotImplementedError
for cls in label.bboxes_name:
if cls not in res:
res.append(cls)
return res
def create_info_file_waymo_wk_fn(idx):
global g_data_dir, g_infos
root_dir = g_data_dir
infos = g_infos
info = dict()
tag = idx
pc_list = [
itm for itm in os.listdir(root_dir) if itm.split('_')[0] == 'velo'
]
pc_paths = {
itm: os.path.join(root_dir, itm, idx + '.npy')
for itm in pc_list
}
img_path = os.path.join(root_dir, "image_front", idx + '.png')
calib = WaymoCalib(os.path.join(root_dir, "calib",
idx + ".txt")).read_calib_file()
label = WaymoLabel(os.path.join(root_dir, "label_imu",
idx + ".txt")).read_label_file()
info["tag"] = tag
info["pc_paths"] = pc_paths
info["img_path"] = img_path
info["calib"] = calib
info["label"] = label
infos.append(info)
def tr_obj(self, label: WaymoLabel, pc: np.array, calib: WaymoCalib,
dx_range: List[float], dy_range: List[float],
dz_range: List[float]) -> (WaymoLabel, np.array):
'''
translate object in the IMU frame
inputs:
label: gt
pc: [#pts, >= 3] in imu frame
calib:
dx_range: [dx_min, dx_max] in imu frame
dy_range: [dy_min, dy_max] in imu frame
dz_range: [dz_min, dz_max] in imu frame
returns:
label_tr
pc_tr
Note: If the attemp times is over max_attemp, it will return the original copy
'''
assert istype(label, "WaymoLabel")
dx_min, dx_max = dx_range
dy_min, dy_max = dy_range
dz_min, dz_max = dz_range
max_attemp = 10
num_attemp = 0
calib_is_iter = isinstance(calib, list)
while True:
num_attemp += 1
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = WaymoLabel()
label_.data = []
for obj in label.data:
label_.data.append(WaymoObj(str(obj)))
if num_attemp > max_attemp:
print(
"WaymoAugmentor.tr_obj: Warning: the attemp times is over {} times,"
"will return the original copy".format(max_attemp))
break
if not calib_is_iter:
for obj in label_.data:
dx = np.random.rand() * (dx_max - dx_min) + dx_min
dy = np.random.rand() * (dy_max - dy_min) + dy_min
dz = np.random.rand() * (dz_max - dz_min) + dz_min
# modify pc
dtr = np.array([dx, dy, dz]).reshape(1, -1)
if isinstance(pc_, dict):
for velo in pc_.keys():
idx = obj.get_pts_idx(pc_[velo][:, :3], calib)
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], dtr)
else:
idx = obj.get_pts_idx(pc_[:, :3], calib)
pc_[idx, :3] = apply_tr(pc_[idx, :3], dtr)
# modify obj
obj.x += dx
obj.y += dy
obj.z += dz
else:
for obj, calib_ in zip(label_.data, calib):
dx = np.random.rand() * (dx_max - dx_min) + dx_min
dy = np.random.rand() * (dy_max - dy_min) + dy_min
dz = np.random.rand() * (dz_max - dz_min) + dz_min
# modify pc
dtr = np.array([dx, dy, dz]).reshape(1, -1)
if isinstance(pc_, dict):
for velo in pc_.keys():
idx = obj.get_pts_idx(pc_[velo][:, :3], calib_)
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], dtr)
else:
idx = obj.get_pts_idx(pc_[:, :3], calib_)
pc_[idx, :3] = apply_tr(pc_[idx, :3], dtr)
# modify obj
obj.x += dx
obj.y += dy
obj.z += dz
if self.check_overlap(label_):
break
res_label = WaymoLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
def rotate_obj(self, label: WaymoLabel, pc: np.array, calib: WaymoCalib,
dry_range: List[float]) -> (WaymoLabel, np.array):
'''
rotate object along the z axis in the LiDAR frame
inputs:
label: gt
pc: [#pts, >= 3] in IMU frame
calib:
dry_range: [dry_min, dry_max] in radius
returns:
label_rot
pc_rot
Note: If the attemp times is over max_attemp, it will return the original copy
'''
assert istype(label, "WaymoLabel")
dry_min, dry_max = dry_range
max_attemp = 10
num_attemp = 0
calib_is_iter = isinstance(calib, list)
while True:
num_attemp += 1
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = WaymoLabel()
label_.data = []
for obj in label.data:
label_.data.append(WaymoObj(str(obj)))
if num_attemp > max_attemp:
print(
"WaymoAugmentor.rotate_obj: Warning: the attemp times is over {} times,"
"will return the original copy".format(max_attemp))
break
if not calib_is_iter:
for obj in label_.data:
dry = np.random.rand() * (dry_max - dry_min) + dry_min
# modify pc
bottom_Fimu = np.array([obj.x, obj.y,
obj.z]).reshape(1, -1)
if isinstance(pc_, dict):
for velo in pc_.keys():
idx = obj.get_pts_idx(pc_[velo][:, :3], calib)
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], -bottom_Fimu)
pc_[velo][idx, :3] = apply_R(
pc_[velo][idx, :3], rotz(dry))
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], bottom_Fimu)
else:
idx = obj.get_pts_idx(pc_[:, :3], calib)
pc_[idx, :3] = apply_tr(pc_[idx, :3], -bottom_Fimu)
pc_[idx, :3] = apply_R(pc_[idx, :3], rotz(dry))
pc_[idx, :3] = apply_tr(pc_[idx, :3], bottom_Fimu)
# modify obj
obj.ry += dry
else:
for obj, calib_ in zip(label_.data, calib):
dry = np.random.rand() * (dry_max - dry_min) + dry_min
# modify pc
bottom_Fimu = np.array([obj.x, obj.y,
obj.z]).reshape(1, -1)
if isinstance(pc_, dict):
for velo in pc_.keys():
idx = obj.get_pts_idx(pc_[velo][:, :3], calib_)
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], -bottom_Fimu)
pc_[velo][idx, :3] = apply_R(
pc_[velo][idx, :3], rotz(dry))
pc_[velo][idx, :3] = apply_tr(
pc_[velo][idx, :3], bottom_Fimu)
else:
idx = obj.get_pts_idx(pc_[:, :3], calib_)
pc_[idx, :3] = apply_tr(pc_[idx, :3], -bottom_Fimu)
pc_[idx, :3] = apply_R(pc_[idx, :3], rotz(dry))
pc_[idx, :3] = apply_tr(pc_[idx, :3], bottom_Fimu)
# modify obj
obj.ry += dry
if self.check_overlap(label_):
break
res_label = WaymoLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_