def test_label_object(data_loader: ArgoverseTrackingLoader) -> None:
label_at_frame_0 = data_loader.get_label_object(0)
for label in label_at_frame_0:
assert label.label_class == "VEHICLE"
assert label.width == 2
assert label.height == 2
assert label.length == 2
def shuffle_log(subset, log: ArgoverseTrackingLoader):
index = np.arange(log.num_lidar_frame)
random.shuffle(index)
for idx in index:
lidar = log.get_lidar(idx)
label = log.get_label_object(idx)
yield idx, subset, lidar, label, log
def extract_datapoints(root_dir, test_set=False):
argoverse_loader = ArgoverseTrackingLoader(root_dir=root_dir)
data = []
for log_id in argoverse_loader.log_list:
# print(f"Processing log: {os.path.join(root_dir, log_id)}", end="\r")
argoverse_data = argoverse_loader.get(log_id=log_id)
# calibL = argoverse_data.get_calibration(camera=camL, log_id=log_id)
# calibR = argoverse_data.get_calibration(camera=camR, log_id=log_id)
calibs = {}
for cam in STEREO_CAMERA_LIST + RING_CAMERA_LIST:
calibs[cam] = argoverse_data.get_calibration(camera=cam,
log_id=log_id)
count = 0
for lidar_timestamp in argoverse_data.lidar_timestamp_list:
data_point = dict()
data_point["log_id"] = log_id
data_point["frame_id"] = count
count += 1
for cam in STEREO_CAMERA_LIST + RING_CAMERA_LIST:
cam_timestamp = argoverse_loader.sync.get_closest_cam_channel_timestamp(
lidar_timestamp=lidar_timestamp,
camera_name=cam,
log_id=log_id)
if cam_timestamp is not None:
data_point[cam] = argoverse_loader.get_image_at_timestamp(
timestamp=cam_timestamp,
camera=cam,
log_id=log_id,
load=False)
else:
data_point[cam] = None
data_point["timestamp"] = lidar_timestamp
data_point["lidar"] = argoverse_loader.timestamp_lidar_dict[
lidar_timestamp]
data_point["calibs"] = calibs
d = argoverse_data.get_pose(
argoverse_data.get_idx_from_timestamp(
lidar_timestamp)).translation
r = Rotation.from_dcm(
argoverse_data.get_pose(
argoverse_data.get_idx_from_timestamp(
lidar_timestamp)).rotation)
data_point["pose"] = (d, r.as_euler('xyz'))
if not test_set:
data_point["labels"] = argoverse_loader.get_label_object(
idx=argoverse_loader.lidar_timestamp_list.index(
lidar_timestamp),
log_id=log_id)
data.append(data_point)
return data
def draw_point_cloud(
ax: plt.Axes,
title: str,
argoverse_data: ArgoverseTrackingLoader,
idx: int,
axes: Optional[Any] = None,
xlim3d: Any = None,
ylim3d: Any = None,
zlim3d: Any = None,
) -> None:
axes = _get_axes_or_default(axes)
pc = argoverse_data.get_lidar(idx)
assert isinstance(pc, np.ndarray)
objects = argoverse_data.get_label_object(idx)
ax.scatter(*np.transpose(pc[:, axes]),
s=point_size,
c=pc[:, 2],
cmap="gray")
ax.set_title(title)
ax.set_xlabel("{} axis".format(axes_str[axes[0]]))
ax.set_ylabel("{} axis".format(axes_str[axes[1]]))
if len(axes) > 2:
ax.set_xlim3d(*axes_limits[axes[0]])
ax.set_ylim3d(*axes_limits[axes[1]])
ax.set_zlim3d(*axes_limits[axes[2]])
ax.set_zlabel("{} axis".format(axes_str[axes[2]]))
else:
ax.set_xlim(*axes_limits[axes[0]])
ax.set_ylim(*axes_limits[axes[1]])
# User specified limits
if xlim3d != None:
ax.set_xlim3d(xlim3d)
if ylim3d != None:
ax.set_ylim3d(ylim3d)
if zlim3d != None:
ax.set_zlim3d(zlim3d)
for obj in objects:
if obj.occlusion == 100:
continue
if obj.label_class is None:
logger.warning("No label class, just picking the default color.")
color = _COLOR_MAP[-1]
else:
color = _COLOR_MAP[OBJ_CLASS_MAPPING_DICT[obj.label_class]]
draw_box(ax, obj.as_3d_bbox().T, axes=axes, color=color)
def draw_point_cloud_trajectory(
ax: plt.Axes,
title: str,
argoverse_data: ArgoverseTrackingLoader,
idx: int,
axes: Optional[Any] = None,
xlim3d: Any = None,
ylim3d: Any = None,
zlim3d: Any = None,
) -> None:
axes = _get_axes_or_default(axes)
unique_id_list = set()
for i in range(len(argoverse_data.lidar_list)):
for label in argoverse_data.get_label_object(i):
unique_id_list.add(label.track_id)
color_map = {
track_id: (
float(np.random.rand()),
float(np.random.rand()),
float(np.random.rand()),
)
for track_id in unique_id_list
}
pc = argoverse_data.get_lidar(idx)
assert isinstance(pc, np.ndarray)
objects = argoverse_data.get_label_object(idx)
ax.scatter(*np.transpose(pc[:, axes]),
s=point_size,
c=pc[:, 2],
cmap="gray")
ax.set_title(title)
ax.set_xlabel("{} axis".format(axes_str[axes[0]]))
ax.set_ylabel("{} axis".format(axes_str[axes[1]]))
if len(axes) > 2:
ax.set_xlim3d(*axes_limits[axes[0]])
ax.set_ylim3d(*axes_limits[axes[1]])
ax.set_zlim3d(*axes_limits[axes[2]])
ax.set_zlabel("{} axis".format(axes_str[axes[2]]))
else:
ax.set_xlim(*axes_limits[axes[0]])
ax.set_ylim(*axes_limits[axes[1]])
# User specified limits
if xlim3d != None:
ax.set_xlim3d(xlim3d)
if ylim3d != None:
ax.set_ylim3d(ylim3d)
if zlim3d != None:
ax.set_zlim3d(zlim3d)
visible_track_id = set()
for obj in objects:
if obj.occlusion == 100:
continue
draw_box(ax,
obj.as_3d_bbox().T,
axes=axes,
color=color_map[obj.track_id])
visible_track_id.add(obj.track_id)
current_pose = argoverse_data.get_pose(idx)
traj_by_id: Dict[Optional[str], List[Any]] = defaultdict(list)
for i in range(0, idx, 1):
if current_pose is None:
logger.warning("`current_pose` is missing at index %d", idx)
break
pose = argoverse_data.get_pose(i)
if pose is None:
logger.warning("`pose` is missing at index %d", i)
continue
objects = argoverse_data.get_label_object(i)
for obj in objects:
if obj.occlusion == 100:
continue
if obj.track_id is None or obj.track_id not in visible_track_id:
continue
x, y, z = pose.transform_point_cloud(
np.array([np.array(obj.translation)]))[0]
x, y, _ = current_pose.inverse_transform_point_cloud(
np.array([np.array([x, y, z])]))[0]
# ax.scatter(x,y, s=point_size, c=color_map[obj.track_id])
if obj.track_id is None:
logger.warning(
"Label has no track_id. Collisions with other tracks that are missing IDs could happen"
)
traj_by_id[obj.track_id].append([x, y])
for track_id in traj_by_id.keys():
traj = np.array(traj_by_id[track_id])
ax.plot(
traj[:, 0],
traj[:, 1],
color=color_map[track_id],
linestyle="--",
linewidth=1,
)
def make_grid_ring_camera(argoverse_data: ArgoverseTrackingLoader,
idx: int) -> Tuple[plt.Figure, plt.Axes]:
f, ax = plt.subplots(3, 3, figsize=(20, 15))
camera = "ring_front_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 0].imshow(img_vis)
ax[0, 0].set_title("Ring Front Left")
ax[0, 0].axis("off")
camera = "ring_front_center"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 1].imshow(img_vis)
ax[0, 1].set_title("Right Front Center")
ax[0, 1].axis("off")
camera = "ring_front_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[0, 2].imshow(img_vis)
ax[0, 2].set_title("Ring Front Right")
ax[0, 2].axis("off")
camera = "ring_side_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[1, 0].imshow(img_vis)
ax[1, 0].set_title("Ring Side Left")
ax[1, 0].axis("off")
ax[1, 1].axis("off")
camera = "ring_side_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[1, 2].imshow(img_vis)
ax[1, 2].set_title("Ring Side Right")
ax[1, 2].axis("off")
camera = "ring_rear_left"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[2, 0].imshow(img_vis)
ax[2, 0].set_title("Ring Rear Left")
ax[2, 0].axis("off")
ax[2, 1].axis("off")
camera = "ring_rear_right"
img = argoverse_data.get_image_sync(idx, camera=camera)
objects = argoverse_data.get_label_object(idx)
calib = argoverse_data.get_calibration(camera)
img_vis = Image.fromarray(show_image_with_boxes(img, objects, calib))
ax[2, 2].imshow(img_vis)
ax[2, 2].set_title("Ring Rear Right")
ax[2, 2].axis("off")
return f, ax
