def test_get_log_calibration_data(
data_loader: SimpleArgoverseTrackingDataLoader, ) -> None:
# Just check that it doesn't raise.
assert data_loader.get_log_calibration_data(_LOG_ID)
def test_get_city_SE3_egovehicle(
data_loader: SimpleArgoverseTrackingDataLoader, ) -> None:
assert data_loader.get_city_SE3_egovehicle(_LOG_ID, 0) is not None
assert data_loader.get_city_SE3_egovehicle(_LOG_ID, 100) is None
def data_loader() -> SimpleArgoverseTrackingDataLoader:
return SimpleArgoverseTrackingDataLoader(os.fspath(_TEST_DATA),
os.fspath(_TEST_DATA))
def test_get_city_name(data_loader: SimpleArgoverseTrackingDataLoader) -> None:
assert data_loader.get_city_name(_LOG_ID) == "PIT"
def dump_clipped_3d_cuboids_to_images(
log_ids: Sequence[str],
max_num_images_to_render: int,
data_dir: str,
experiment_prefix: str,
motion_compensate: bool = True,
omit_centerlines: bool = False,
generate_video_only: bool = False,
) -> List[str]:
"""
We bring the 3D points into each camera coordinate system, and do the clipping there in 3D.
Args:
log_ids: A list of log IDs
max_num_images_to_render: maximum numbers of images to render.
data_dir: path to dataset with the latest data
experiment_prefix: Output directory
motion_compensate: Whether to motion compensate when projecting
omit_centerlines: whether to omit map vector lane geometry from rendering
generate_video_only: whether to generate mp4s only without dumping individual frames
Returns:
saved_img_fpaths
"""
saved_img_fpaths = []
dl = SimpleArgoverseTrackingDataLoader(data_dir=data_dir,
labels_dir=data_dir)
if not omit_centerlines:
avm = ArgoverseMap()
fps_map = {
cam_name: STEREO_FPS if "stereo" in cam_name else RING_CAM_FPS
for cam_name in RING_CAMERA_LIST + STEREO_CAMERA_LIST
}
category_subdir = "amodal_labels"
if not Path(f"{experiment_prefix}_{category_subdir}").exists():
os.makedirs(f"{experiment_prefix}_{category_subdir}")
video_output_dir = f"{experiment_prefix}_{category_subdir}"
for log_id in log_ids:
save_dir = f"{experiment_prefix}_{log_id}"
if not generate_video_only and not Path(save_dir).exists():
# JPG images will be dumped here, if requested by arguments
os.makedirs(save_dir)
city_name = dl.get_city_name(log_id)
log_calib_data = dl.get_log_calibration_data(log_id)
flag_done = False
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST +
STEREO_CAMERA_LIST):
fps = fps_map[camera_name]
if generate_video_only:
mp4_path = f"{video_output_dir}/{log_id}_{camera_name}_{fps}fps.mp4"
video_writer = VideoWriter(mp4_path)
cam_im_fpaths = dl.get_ordered_log_cam_fpaths(log_id, camera_name)
for i, im_fpath in enumerate(cam_im_fpaths):
if i % 50 == 0:
logging.info("\tOn file %s of camera %s of %s", i,
camera_name, log_id)
cam_timestamp = Path(im_fpath).stem.split("_")[-1]
cam_timestamp = int(cam_timestamp)
# load PLY file path, e.g. 'PC_315978406032859416.ply'
ply_fpath = dl.get_closest_lidar_fpath(log_id, cam_timestamp)
if ply_fpath is None:
continue
lidar_pts = load_ply(ply_fpath)
save_img_fpath = f"{save_dir}/{camera_name}_{cam_timestamp}.jpg"
if Path(save_img_fpath).exists():
saved_img_fpaths += [save_img_fpath]
if max_num_images_to_render != -1 and len(
saved_img_fpaths) > max_num_images_to_render:
flag_done = True
break
continue
city_SE3_egovehicle = dl.get_city_SE3_egovehicle(
log_id, cam_timestamp)
if city_SE3_egovehicle is None:
continue
lidar_timestamp = Path(ply_fpath).stem.split("_")[-1]
lidar_timestamp = int(lidar_timestamp)
labels = dl.get_labels_at_lidar_timestamp(
log_id, lidar_timestamp)
if labels is None:
logging.info("\tLabels missing at t=%s", lidar_timestamp)
continue
# Swap channel order as OpenCV expects it -- BGR not RGB
# must make a copy to make memory contiguous
img = imageio.imread(im_fpath)[:, :, ::-1].copy()
camera_config = get_calibration_config(log_calib_data,
camera_name)
planes = generate_frustum_planes(
camera_config.intrinsic.copy(), camera_name)
if not omit_centerlines:
img = plot_lane_centerlines_in_img(
lidar_pts,
city_SE3_egovehicle,
img,
city_name,
avm,
camera_config,
planes,
)
for label_idx, label in enumerate(labels):
obj_rec = json_label_dict_to_obj_record(label)
if obj_rec.occlusion == 100:
continue
cuboid_vertices = obj_rec.as_3d_bbox()
points_h = point_cloud_to_homogeneous(cuboid_vertices).T
if motion_compensate:
(
uv,
uv_cam,
valid_pts_bool,
K,
) = project_lidar_to_img_motion_compensated(
points_h, # these are recorded at lidar_time
copy.deepcopy(log_calib_data),
camera_name,
cam_timestamp,
lidar_timestamp,
data_dir,
log_id,
return_K=True,
)
else:
# project_lidar_to_img
(
uv,
uv_cam,
valid_pts_bool,
camera_config,
) = project_lidar_to_undistorted_img(
points_h, copy.deepcopy(log_calib_data),
camera_name)
if valid_pts_bool.sum() == 0:
continue
img = obj_rec.render_clip_frustum_cv2(
img,
uv_cam.T[:, :3],
planes.copy(),
copy.deepcopy(camera_config),
)
if generate_video_only:
video_writer.add_frame(rgb_frame=img[:, :, ::-1])
else:
cv2.imwrite(save_img_fpath, img)
saved_img_fpaths += [save_img_fpath]
if (not generate_video_only and max_num_images_to_render != -1
and len(saved_img_fpaths) > max_num_images_to_render):
flag_done = True
break
if generate_video_only:
video_writer.complete()
ffmpeg_compress_video(mp4_path, fps)
if flag_done:
break
if not generate_video_only:
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST +
STEREO_CAMERA_LIST):
# Write the cuboid video from individual frames -- could also write w/ fps=20,30,40
fps = fps_map[camera_name]
img_wildcard = f"{save_dir}/{camera_name}_%*.jpg"
output_fpath = f"{video_output_dir}/{log_id}_{camera_name}_{fps}fps.mp4"
write_nonsequential_idx_video(img_wildcard, output_fpath, fps)
return saved_img_fpaths
def dump_clipped_3d_cuboids_to_images(
log_ids: Sequence[str],
max_num_images_to_render: int,
data_dir: str,
experiment_prefix: str,
is_shrinkwrap: bool,
motion_compensate: bool = False,
) -> List[str]:
"""
We bring the 3D points into each camera coordinate system, and do the clipping there in 3D.
Args:
log_ids: A list of log IDs
max_num_images_to_render: maximum numbers of images to render.
data_dir: path to dataset with the latest data
experiment_prefix: Output directory
is_shrinkwrap: Whether to place in a `"shrinkwrap"` output directory
motion_compensate: Whether to motion compensate when projecting
Returns:
saved_img_fpaths
"""
saved_img_fpaths = []
dl = SimpleArgoverseTrackingDataLoader(data_dir=data_dir,
labels_dir=data_dir)
avm = ArgoverseMap()
for log_id in log_ids:
save_dir = f"{experiment_prefix}_{log_id}"
if not Path(save_dir).exists():
os.makedirs(save_dir)
city_name = dl.get_city_name(log_id)
log_calib_data = dl.get_log_calibration_data(log_id)
ply_fpaths = dl.get_ordered_log_ply_fpaths(log_id)
flag_done = False
for i, ply_fpath in enumerate(ply_fpaths):
if i % 50 == 0:
logging.info(f"\tOn file {i} of {log_id}")
lidar_timestamp = ply_fpath.split("/")[-1].split(".")[0].split(
"_")[-1]
lidar_timestamp = int(lidar_timestamp)
lidar_pts = load_ply(ply_fpath)
city_to_egovehicle_se3 = dl.get_city_to_egovehicle_se3(
log_id, lidar_timestamp)
if city_to_egovehicle_se3 is None:
continue
labels = dl.get_labels_at_lidar_timestamp(log_id, lidar_timestamp)
if labels is None:
logger.info(f"\tLabels missing at t={lidar_timestamp}")
continue
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST +
STEREO_CAMERA_LIST):
camera_counter = 0
# load images, e.g. 'image_raw_ring_front_center_000000486.jpg'
im_fpath = dl.get_closest_im_fpath(log_id, camera_name,
lidar_timestamp)
if im_fpath is None:
continue
save_img_fpath = f"{save_dir}/{camera_name}_{lidar_timestamp}.jpg"
if Path(save_img_fpath).exists():
saved_img_fpaths += [save_img_fpath]
if max_num_images_to_render != -1 and len(
saved_img_fpaths) > max_num_images_to_render:
flag_done = True
break
continue
# Swap channel order as OpenCV expects it -- BGR not RGB
# must make a copy to make memory contiguous
img = imageio.imread(im_fpath)[:, :, ::-1].copy()
camera_config = get_calibration_config(log_calib_data,
camera_name)
planes = generate_frustum_planes(
camera_config.intrinsic.copy(), camera_name)
img = plot_lane_centerlines_in_img(lidar_pts,
city_to_egovehicle_se3, img,
city_name, avm,
camera_config, planes)
for label_idx, label in enumerate(labels):
obj_rec = json_label_dict_to_obj_record(label)
if obj_rec.occlusion == 100:
continue
cuboid_vertices = obj_rec.as_3d_bbox()
points_h = point_cloud_to_homogeneous(cuboid_vertices).T
if motion_compensate:
uv, uv_cam, valid_pts_bool, K = project_lidar_to_img_motion_compensated(
points_h, # these are recorded at lidar_time
copy.deepcopy(log_calib_data),
camera_name,
cam_timestamp,
lidar_timestamp,
data_dir,
log_id,
return_K=True,
)
else:
# project_lidar_to_img
uv, uv_cam, valid_pts_bool, camera_config = project_lidar_to_undistorted_img(
points_h, copy.deepcopy(log_calib_data),
camera_name)
if valid_pts_bool.sum() == 0:
continue
img = obj_rec.render_clip_frustum_cv2(
img, uv_cam.T[:, :3], planes.copy(),
copy.deepcopy(camera_config))
cv2.imwrite(save_img_fpath, img)
saved_img_fpaths += [save_img_fpath]
if max_num_images_to_render != -1 and len(
saved_img_fpaths) > max_num_images_to_render:
flag_done = True
break
if flag_done:
break
category_subdir = "amodal_labels"
if not Path(f"{experiment_prefix}_{category_subdir}").exists():
os.makedirs(f"{experiment_prefix}_{category_subdir}")
for cam_idx, camera_name in enumerate(RING_CAMERA_LIST +
STEREO_CAMERA_LIST):
# Write the cuboid video -- could also write w/ fps=20,30,40
if "stereo" in camera_name:
fps = 5
else:
fps = 10
img_wildcard = f"{save_dir}/{camera_name}_%*.jpg"
output_fpath = f"{experiment_prefix}_{category_subdir}/{log_id}_{camera_name}_{fps}fps.mp4"
write_nonsequential_idx_video(img_wildcard, output_fpath, fps)
return saved_img_fpaths
def test_get_labels_at_lidar_timestamp(
data_loader: SimpleArgoverseTrackingDataLoader) -> None:
assert data_loader.get_labels_at_lidar_timestamp(_LOG_ID, 0) is not None
assert data_loader.get_labels_at_lidar_timestamp(_LOG_ID, 100) is None
def test_get_ordered_log_ply_fpaths(
data_loader: SimpleArgoverseTrackingDataLoader) -> None:
# Test data doesn't have cameras so we cannot currently test this if we
assert len(data_loader.get_ordered_log_ply_fpaths(_LOG_ID)) == 3
def test_get_closest_im_fpath(
data_loader: SimpleArgoverseTrackingDataLoader) -> None:
# Test data doesn't have cameras so we cannot currently test this if we
assert data_loader.get_closest_im_fpath(_LOG_ID, "camera_name", 0) is None