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 main():
parser = argparse.ArgumentParser(
description='Split Dataset according to a txt file.')
parser.add_argument('--input-dir',
type=str,
metavar='INPUT DIR',
help='input dir')
parser.add_argument('--calib-dir',
type=str,
metavar='CALIB DIR',
help='calib dir')
parser.add_argument('--output-dir',
type=str,
metavar='OUTPUT DIR',
help='output dir')
args = parser.parse_args()
input_dir = args.input_dir
calib_dir = args.calib_dir
output_dir = args.output_dir
assert os.path.isdir(input_dir)
assert os.path.isdir(calib_dir)
assert os.path.isdir(output_dir)
idx_list = os.listdir(input_dir)
num = len(idx_list)
assert num == len(os.listdir(calib_dir))
for idx in idx_list:
print(idx)
label = CarlaLabel(os.path.join(input_dir, idx)).read_label_file()
calib = CarlaCalib(os.path.join(calib_dir, idx)).read_calib_file()
label_Fcam = CarlaLabel()
label_Fcam.data = []
for obj in label.data:
obj_Fcam = CarlaObj(str(obj))
cns_Fimu = obj.get_bbox3dcorners()
cns_Fcam = calib.imu2cam(cns_Fimu)
cns_Fcam2d = calib.cam2imgplane(cns_Fcam)
minx = float(np.min(cns_Fcam2d[:, 0]))
maxx = float(np.max(cns_Fcam2d[:, 0]))
miny = float(np.min(cns_Fcam2d[:, 1]))
maxy = float(np.max(cns_Fcam2d[:, 1]))
obj_Fcam.score = obj.score
obj_Fcam.x, obj_Fcam.y, obj_Fcam.z = calib.imu2cam(
np.array([obj.x, obj.y, obj.z]).reshape(1, 3))[0]
obj_Fcam.w, obj_Fcam.l, obj_Fcam.h = obj.l, obj.h, obj.w
obj_Fcam.ry = -obj.ry
obj_Fcam.bbox_l, obj_Fcam.bbox_t, obj_Fcam.bbox_r, obj_Fcam.bbox_b = minx, miny, maxx, maxy
label_Fcam.data.append(obj_Fcam)
# validate("/usr/app/data/CARLA/train/", idx, label_Fcam, "/usr/app/vis/train/")
write_str_to_file(str(label_Fcam), os.path.join(output_dir, idx))
def flip_pc(self, label: CarlaLabel, pc: (np.array, dict),
calib: CarlaCalib) -> (CarlaLabel, 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, "CarlaLabel")
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = CarlaLabel()
label_.data = []
for obj in label.data:
label_.data.append(CarlaObj(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 = CarlaLabel()
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 keep(self, label: CarlaLabel, pc: (np.array, dict),
calib: CarlaCalib) -> (CarlaLabel, np.array):
assert istype(label, "CarlaLabel")
# copy pc & label
if isinstance(pc, dict):
pc_ = {k: v.copy() for k, v in pc.items()}
else:
pc_ = pc.copy()
label_ = CarlaLabel()
label_.data = []
for obj in label.data:
label_.data.append(CarlaObj(str(obj)))
res_label = CarlaLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
def test_get_pts_idx(self):
from det3.ops import crop_pts_3drot, read_npy
from det3.dataloader.carladata import CarlaObj, CarlaObj, CarlaLabel, CarlaCalib
pts = read_npy("./unit-test/data/test_CarlaAugmentor_000250.npy")
calib = CarlaCalib(
"./unit-test/data/test_CarlaAugmentor_000250_calib.txt"
).read_calib_file()
label = CarlaLabel(
"./unit-test/data/test_CarlaAugmentor_000250_label_imu.txt"
).read_label_file()
pts = calib.lidar2imu(pts, key='Tr_imu_to_velo_top')
for i in range(len(label)):
obj = UdiObj(arr=np.array(label.bboxes3d[i,
[3, 4, 5, 2, 1, 0, 6]]),
cls="Car",
score=0.99)
idx = obj.get_pts_idx(pts)
gt = label.data[i].get_pts_idx(pts)
gt = np.where(gt)[0].flatten()
self.assertTrue(set(idx.tolist()) == set(gt.tolist()))
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 test_corner_transformation(self):
from det3.ops import crop_pts_3drot, read_npy
from det3.dataloader.carladata import CarlaObj, CarlaObj, CarlaLabel, CarlaCalib
pts = read_npy("./unit-test/data/test_CarlaAugmentor_000250.npy")
calib = CarlaCalib(
"./unit-test/data/test_CarlaAugmentor_000250_calib.txt"
).read_calib_file()
label = CarlaLabel(
"./unit-test/data/test_CarlaAugmentor_000250_label_imu.txt"
).read_label_file()
pts = calib.lidar2imu(pts, key='Tr_imu_to_velo_top')
for i in range(len(label)):
obj = UdiObj(arr=np.array(label.bboxes3d[i,
[3, 4, 5, 2, 1, 0, 6]]),
cls="Car",
score=0.99)
cns = obj.get_bbox3d_corners()
gt = label.data[i].get_bbox3dcorners()
self.assertTrue(np.allclose(cns, gt, atol=1e-2))
obj_cp = UdiObj()
obj_cp.from_corners(cns, obj.cls, obj.score)
self.assertTrue(obj_cp.equal(obj))
def create_info_file_carla_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_2", idx + '.png')
calib = CarlaCalib(os.path.join(root_dir, "calib",
idx + ".txt")).read_calib_file()
label = CarlaLabel(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 test_crop_pts_3drot(self):
print("test_crop_pts_3drot")
from det3.ops import crop_pts_3drot, read_npy
from det3.dataloader.carladata import CarlaObj, CarlaObj, CarlaLabel, CarlaCalib
pts = read_npy("./unit-test/data/test_CarlaAugmentor_000250.npy")
calib = CarlaCalib(
"./unit-test/data/test_CarlaAugmentor_000250_calib.txt"
).read_calib_file()
label = CarlaLabel(
"./unit-test/data/test_CarlaAugmentor_000250_label_imu.txt"
).read_label_file()
pts = calib.lidar2imu(pts, key='Tr_imu_to_velo_top')
for test_dtype in [np.float32, np.float64]:
# obj.h, obj.w, obj.l, obj.x, obj.y, obj.z, obj.ry
boxes = label.bboxes3d[:, [3, 4, 5, 2, 1, 0, 6]].astype(test_dtype)
pts = pts.astype(test_dtype)
gt = [obj.get_pts_idx(pts) for obj in label.data]
gt = [np.where(itm)[0].flatten() for itm in gt]
# np
est = crop_pts_3drot(boxes, pts)
self.assertTrue(len(gt) == len(est))
for gt_, est_ in zip(gt, est):
set_gt = set(gt_.tolist())
set_est = set(est_.tolist())
self.assertTrue(set_gt == set_est)
times = 1000
t1 = time.time()
for i in range(times):
est = crop_pts_3drot(boxes, pts)
t = (time.time() - t1) / times * 1000
print(f"np {times}: {t:.3f} ms {est[0].dtype}")
# torch
pts_ts = torch.from_numpy(pts)
boxes_ts = torch.from_numpy(boxes)
est_ts = crop_pts_3drot(boxes_ts, pts_ts)
self.assertTrue(len(gt) == len(est_ts))
for gt_, est_ts_ in zip(gt, est_ts):
set_gt = set(gt_.tolist())
set_est_ts = set(est_ts_.numpy().tolist())
self.assertTrue(set_gt == set_est_ts)
t1 = time.time()
for i in range(times):
est_ts = crop_pts_3drot(boxes_ts, pts_ts)
t = (time.time() - t1) / times * 1000
print(f"torchcpu {times}: {t:.3f} ms {est_ts[0].dtype}")
# torchgpu
pts_tsgpu = torch.from_numpy(pts).cuda()
boxes_tsgpu = torch.from_numpy(boxes).cuda()
est_tsgpu = crop_pts_3drot(boxes_tsgpu, pts_tsgpu)
self.assertTrue(len(gt) == len(est_tsgpu))
for gt_, est_tsgpu_ in zip(gt, est_tsgpu):
set_gt = set(gt_.tolist())
set_est_tsgpu = set(est_tsgpu_.cpu().numpy().tolist())
self.assertTrue(set_gt == set_est_tsgpu)
torch.cuda.synchronize()
t1 = time.time()
for i in range(times):
est_tsgpu = crop_pts_3drot(boxes_tsgpu, pts_tsgpu)
torch.cuda.synchronize()
t = (time.time() - t1) / times * 1000
print(f"torchgpu {times}: {t:.3f} ms {est_tsgpu[0].dtype}")
def test_equal(self):
label1 = CarlaLabel('./unit-test/data/test_CarlaLabel_000000.txt'
).read_label_file(no_dontcare=True)
label2 = CarlaLabel('./unit-test/data/test_CarlaLabel_000000.txt'
).read_label_file(no_dontcare=True)
self.assertTrue(label1.equal(label1, ['Car', 'Van'], rtol=1e-5))
self.assertTrue(label1.equal(label2, ['Car', 'Van'], rtol=1e-5))
self.assertTrue(label2.equal(label1, ['Car', 'Van'], rtol=1e-5))
self.assertTrue(label2.equal(label2, ['Car', 'Van'], rtol=1e-5))
label1 = CarlaLabel('./unit-test/data/test_CarlaLabel_000000.txt'
).read_label_file(no_dontcare=True)
label2 = CarlaLabel('./unit-test/data/test_CarlaLabel_000010.txt'
).read_label_file(no_dontcare=True)
self.assertTrue(not label1.equal(label2, ['Car', 'Van'], rtol=1e-5))
self.assertTrue(not label2.equal(label1, ['Car', 'Van'], rtol=1e-5))
def test_read_label_file(self):
label = CarlaLabel('./unit-test/data/test_CarlaLabel_000000.txt'
).read_label_file(no_dontcare=True)
self.assertEqual(len(label.data), 60)
def test_init(self):
label = CarlaLabel('./unit-test/data/test_CarlaLabel_000000.txt')
self.assertEqual(label.data, None)
def tr_obj(self, label: CarlaLabel, pc: (np.array, dict),
calib: CarlaCalib, dx_range: List[float], dy_range: List[float],
dz_range: List[float]) -> (CarlaLabel, np.array):
'''
translate object in the IMU frame
inputs:
label: gt
pc: [#pts, >= 3] in imu framem, dict
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, "CarlaLabel")
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_ = CarlaLabel()
label_.data = []
for obj in label.data:
label_.data.append(CarlaObj(str(obj)))
if num_attemp > max_attemp:
print(
"CarlaAugmentor.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 = CarlaLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
def rotate_obj(self, label: CarlaLabel, pc: (np.array, dict),
calib: CarlaCalib,
dry_range: List[float]) -> (CarlaLabel, np.array):
'''
rotate object along the z axis in the LiDAR frame
inputs:
label: gt
pc: [#pts, >= 3] in IMU frame, dict
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, "CarlaLabel")
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_ = CarlaLabel()
label_.data = []
for obj in label.data:
label_.data.append(CarlaObj(str(obj)))
if num_attemp > max_attemp:
print(
"CarlaAugmentor.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 = CarlaLabel()
for obj in label_.data:
res_label.add_obj(obj)
res_label.current_frame = "IMU"
return res_label, pc_
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
