def get_split_cam_configs_and_sync_data(argoverse_tracking_root_dir, split_dir, camera_list, source_views_indexes): """ Returns a dictionary giving the intrinsics/extrinsics/resolution of given cameras for each sequence. Parameters ---------- split_dir: pathlib Path argoverse-tracking split root dir, e.g.: path_to_data/argoverse-tracking/train1 camera_list: list list of camera for which to load the parameters. Must be in Argoverse's CAMERA_LIST Returns ------- dict A dictionary where the key is the log string and the value a dict corresponding to the cameras' parameters """ camera_configs = {} synchronized_data = [] db = SynchronizationDB(str(split_dir)) valid_logs = list(db.get_valid_logs()) # log_ids founds under split_dir split = split_dir.stem # e.g., train1 for log in valid_logs: camera_configs[log] = load_camera_config(str(split_dir / log / VEHICLE_CALIBRATION_INFO_FILENAME), camera_list) synchronized_data += get_synchronized_data(argoverse_tracking_root_dir, db, split, log, camera_list, source_views_indexes) return camera_configs, synchronized_data
def __init__(self, data_dir: str, labels_dir: str) -> None: """ Args: data_dir: str, representing path to raw Argoverse data labels_dir: strrepresenting path to Argoverse data labels """ self.data_dir = data_dir self.labels_dir = labels_dir self.sdb = SynchronizationDB(data_dir)
def __init__(self, data_dir: str, labels_dir: str) -> None: """ Args: data_dir: str, representing path to raw Argoverse data labels_dir: str representing path to Argoverse data labels (e.g. labels or estimated detections/tracks) """ self.data_dir = data_dir self.labels_dir = labels_dir self.sdb = SynchronizationDB(data_dir)
def __init__( self, dataset_dir: str, experiment_prefix: str, use_existing_files: bool = True, log_id: str = None ) -> None: """We will cache the accumulated trajectories per city, per log, and per frame for the tracking benchmark. """ self.plot_lane_tangent_arrows = True self.plot_lidar_bev = True self.plot_lidar_in_img = False self.experiment_prefix = experiment_prefix self.dataset_dir = dataset_dir self.labels_dir = dataset_dir self.sdb = SynchronizationDB(self.dataset_dir) if log_id is None: tmp_dir = tempfile.gettempdir() per_city_traj_dict_fpath = f"{tmp_dir}/per_city_traj_dict_{experiment_prefix}.pkl" log_egopose_dict_fpath = f"{tmp_dir}/log_egopose_dict_{experiment_prefix}.pkl" log_timestamp_dict_fpath = f"{tmp_dir}/log_timestamp_dict_{experiment_prefix}.pkl" if not use_existing_files: # write the accumulate data dictionaries to disk PerFrameLabelAccumulator(dataset_dir, dataset_dir, experiment_prefix) self.per_city_traj_dict = load_pkl_dictionary(per_city_traj_dict_fpath) self.log_egopose_dict = load_pkl_dictionary(log_egopose_dict_fpath) self.log_timestamp_dict = load_pkl_dictionary(log_timestamp_dict_fpath) else: pfa = PerFrameLabelAccumulator(dataset_dir, dataset_dir, experiment_prefix, save=False) pfa.accumulate_per_log_data(log_id=log_id) self.per_city_traj_dict = pfa.per_city_traj_dict self.log_egopose_dict = pfa.log_egopose_dict self.log_timestamp_dict = pfa.log_timestamp_dict
def visualize_ground_lidar_pts(log_id: str, dataset_dir: str, experiment_prefix: str): """Process a log by drawing the LiDAR returns that are classified as belonging to the ground surface in a red to green colormap in the image. Args: log_id: The ID of a log dataset_dir: Where the dataset is stored experiment_prefix: Output prefix """ sdb = SynchronizationDB(dataset_dir, collect_single_log_id=log_id) city_info_fpath = f"{dataset_dir}/{log_id}/city_info.json" city_info = read_json_file(city_info_fpath) city_name = city_info["city_name"] avm = ArgoverseMap() ply_fpaths = sorted(glob.glob(f"{dataset_dir}/{log_id}/lidar/PC_*.ply")) for i, ply_fpath in enumerate(ply_fpaths): if i % 500 == 0: print(f"\tOn file {i} of {log_id}") lidar_timestamp_ns = ply_fpath.split("/")[-1].split(".")[0].split( "_")[-1] pose_fpath = f"{dataset_dir}/{log_id}/poses/city_SE3_egovehicle_{lidar_timestamp_ns}.json" if not Path(pose_fpath).exists(): continue pose_data = read_json_file(pose_fpath) rotation = np.array(pose_data["rotation"]) translation = np.array(pose_data["translation"]) city_to_egovehicle_se3 = SE3(rotation=quat2rotmat(rotation), translation=translation) lidar_pts = load_ply(ply_fpath) lidar_timestamp_ns = int(lidar_timestamp_ns) draw_ground_pts_in_image( sdb, lidar_pts, city_to_egovehicle_se3, avm, log_id, lidar_timestamp_ns, city_name, dataset_dir, experiment_prefix, ) for camera_name in CAMERA_LIST: if "stereo" in camera_name: fps = 5 else: fps = 10 cmd = f"ffmpeg -r {fps} -f image2 -i '{experiment_prefix}_ground_viz/{log_id}/{camera_name}/%*.jpg' {experiment_prefix}_ground_viz/{experiment_prefix}_{log_id}_{camera_name}_{fps}fps.mp4" print(cmd) run_command(cmd)
def __init__( self, dataset_dir: str, labels_dir: str, experiment_prefix: str, bboxes_3d: bool = False, save: bool = True, ) -> None: """Initialize PerFrameLabelAccumulator object for use with tracking benchmark data. Args: dataset_dir (str): Dataset directory. labels_dir (str): Labels directory. experiment_prefix (str): Prefix for experimint to use. bboxes_3d (bool, optional): to use 3d bounding boxes (True) or 2d bounding boxes (False). """ self.bboxes_3d = bboxes_3d self.dataset_dir = dataset_dir self.labels_dir = labels_dir tmp_dir = tempfile.gettempdir() per_city_traj_dict_fpath = f"{tmp_dir}/per_city_traj_dict_{experiment_prefix}.pkl" log_egopose_dict_fpath = f"{tmp_dir}/log_egopose_dict_{experiment_prefix}.pkl" log_timestamp_dict_fpath = f"{tmp_dir}/log_timestamp_dict_{experiment_prefix}.pkl" # coordinate system is the map world frame self.per_city_traj_dict: Dict[str, List[Tuple[np.ndarray, str]]] = { "MIA": [], "PIT": [], } # all the trajectories for these 2 cities self.log_egopose_dict: Dict[str, Dict[int, Dict[str, np.ndarray]]] = {} self.log_timestamp_dict: Dict[str, Dict[int, List[FrameRecord]]] = {} self.sdb = SynchronizationDB(self.dataset_dir) if save: self.accumulate_per_log_data() save_pkl_dictionary(per_city_traj_dict_fpath, self.per_city_traj_dict) save_pkl_dictionary(log_egopose_dict_fpath, self.log_egopose_dict) save_pkl_dictionary(log_timestamp_dict_fpath, self.log_timestamp_dict)
def __init__(self, root_dir: str) -> None: # initialize class member self.CAMERA_LIST = CAMERA_LIST self._log_list: Optional[List[str]] = [None] self._image_list: Optional[Dict[str, Dict[str, List[str]]]] = None self._image_list_sync: Optional[Dict[str, Dict[str, List[np.ndarray]]]] = None self._lidar_list: Optional[Dict[str, List[str]]] = None self._image_timestamp_list: Optional[Dict[str, Dict[str, List[int]]]] = None self._timestamp_image_dict: Optional[Dict[str, Dict[str, Dict[int, str]]]] = None self._image_timestamp_list_sync: Optional[Dict[str, Dict[str, List[int]]]] = None self._lidar_timestamp_list: Optional[Dict[str, List[int]]] = None self._timestamp_lidar_dict: Optional[Dict[str, Dict[int, str]]] = None self._label_list: Optional[Dict[str, List[str]]] = None self._calib: Optional[Dict[str, Dict[ str, Calibration]]] = None # { log_name: { camera_name: Calibration } } self._city_name = None self.counter: int = 0 self.image_count: int = 0 self.lidar_count: int = 0 self.root_dir: str = root_dir self.current_log = self.log_list[self.counter] assert self.image_list is not None assert self.lidar_list is not None assert self.label_list is not None # load calibration file self.calib_filename: str = os.path.join( self.root_dir, self.current_log, "vehicle_calibration_info.json") # lidar @10hz, ring camera @30hz, stereo camera @5hz self.num_lidar_frame: int = len(self.lidar_timestamp_list) self.num_ring_camera_frame: int = len( self.image_timestamp_list[RING_CAMERA_LIST[0]]) self.num_stereo_camera_frame: int = len( self.image_timestamp_list[STEREO_CAMERA_LIST[0]]) self.sync: SynchronizationDB = SynchronizationDB(root_dir) assert self.image_list_sync is not None assert self.calib is not None
def __init__(self, dataset_dir: str, experiment_prefix: str, tf_record_prefix: str) -> None: """We will cache the accumulated trajectories per city, per log, and per frame for the tracking benchmark. """ self.experiment_prefix = experiment_prefix self.tf_record_prefix = tf_record_prefix self.dataset_dir = dataset_dir self.labels_dir = dataset_dir self.sdb = SynchronizationDB(self.dataset_dir) self.per_city_traj_dict = None self.log_egopose_dict = None self.log_timestamp_dict = None self.timestamps = None
def main(argo_tracking_root_dir, output_dir, cameras, acc_sweeps, ip_basic): print('Preprocessing data....') print("INPUT DIR: ", argo_tracking_root_dir) print("OUTPUT DIR: ", output_dir) if cameras is None: cameras = CAMERA_LIST else: cameras = validate_camera_option(cameras) print(cameras) argo_tracking_root_dir = Path(argo_tracking_root_dir).expanduser() output_dir = Path(output_dir).expanduser() split_namedirs = ["train1", "train2", "train3", "train4", "test", "val"] split_dirs = [ argo_tracking_root_dir / split_namedir for split_namedir in split_namedirs ] for split_dir in split_dirs: db = SynchronizationDB(str(split_dir)) log_dirs = sorted(list(split_dir.iterdir())) for log_dir in log_dirs: lidar_dir = log_dir / "lidar" lidar_filepaths = [f for f in sorted(list(lidar_dir.iterdir()))] total = len(lidar_filepaths) calib_filepath = str(log_dir / "vehicle_calibration_info.json") output_base_path = output_dir / split_dir.stem / log_dir.stem with open(calib_filepath, "r") as f: calib_data = json.load(f) args = (output_base_path, calib_data, cameras, db, acc_sweeps, ip_basic) with tqdm(lidar_filepaths, desc=f"{split_dir.stem} | log {log_dir.stem} ", total=total) as progress: for lidar_filepath in progress: project_and_save(lidar_filepath, *args) print('Preprocessing of LiDAR data Finished.')
def generate_vehicle_bev(dataset_dir="", log_id="", output_dir=""): argoverse_loader = ArgoverseTrackingLoader(dataset_dir) argoverse_data = argoverse_loader.get(log_id) camera = argoverse_loader.CAMERA_LIST[7] calib = argoverse_data.get_calibration(camera) sdb = SynchronizationDB(dataset_dir, collect_single_log_id=log_id) calib_path = f"{dataset_dir}/{log_id}/vehicle_calibration_info.json" calib_data = read_json_file(calib_path) ind = 0 for i in range(9): if calib_data['camera_data_'][i]['key'] == \ 'image_raw_stereo_front_left': ind = i break rotation = np.array(calib_data['camera_data_'][ind]['value'] ['vehicle_SE3_camera_']['rotation']['coefficients']) translation = np.array(calib_data['camera_data_'][ind]['value'] ['vehicle_SE3_camera_']['translation']) egovehicle_SE3_cam = SE3(rotation=quat2rotmat(rotation), translation=translation) if not os.path.exists(os.path.join(output_dir, log_id, "car_bev_gt")): os.makedirs(os.path.join(output_dir, log_id, "car_bev_gt")) lidar_dir = os.path.join(dataset_dir, log_id, "lidar") ply_list = os.listdir(lidar_dir) pfa = PerFrameLabelAccumulator(dataset_dir, dataset_dir, "argoverse_bev_viz", save=False, bboxes_3d=True) pfa.accumulate_per_log_data(log_id) log_timestamp_dict = pfa.log_timestamp_dict for i, ply_name in enumerate(ply_list): lidar_timestamp = ply_name.split('.')[0].split('_')[1] lidar_timestamp = int(lidar_timestamp) cam_timestamp = sdb.get_closest_cam_channel_timestamp( lidar_timestamp, "stereo_front_left", str(log_id)) image_path = os.path.join( output_dir, str(log_id), "car_bev_gt", "stereo_front_left_" + str(cam_timestamp) + ".jpg") objects = log_timestamp_dict[log_id][lidar_timestamp] top_view = np.zeros((256, 256)) all_occluded = True for frame_rec in objects: if frame_rec.occlusion_val != IS_OCCLUDED_FLAG: all_occluded = False if not all_occluded: for i, frame_rec in enumerate(objects): bbox_ego_frame = frame_rec.bbox_ego_frame uv = calib.project_ego_to_image(bbox_ego_frame).T idx_ = np.all( np.logical_and( np.logical_and( np.logical_and(uv[0, :] >= 0.0, uv[0, :] < size[1] - 1.0), np.logical_and(uv[1, :] >= 0.0, uv[1, :] < size[0] - 1.0)), uv[2, :] > 0)) if not idx_: continue bbox_cam_fr = egovehicle_SE3_cam.inverse().\ transform_point_cloud(bbox_ego_frame) X = bbox_cam_fr[:, 0] Z = bbox_cam_fr[:, 2] if (frame_rec.occlusion_val != IS_OCCLUDED_FLAG and frame_rec.obj_class_str == "VEHICLE"): y_img = (-Z / res).astype(np.int32) x_img = (X / res).astype(np.int32) x_img -= int(np.floor(-20 / res)) y_img += int(np.floor(40 / res)) box = np.array([x_img[2], y_img[2]]) box = np.vstack((box, [x_img[6], y_img[6]])) box = np.vstack((box, [x_img[7], y_img[7]])) box = np.vstack((box, [x_img[3], y_img[3]])) cv2.drawContours(top_view, [box], 0, 255, -1) cv2.imwrite(image_path, top_view)
def generate_road_bev(dataset_dir="", log_id="", output_dir=""): argoverse_loader = ArgoverseTrackingLoader(dataset_dir) argoverse_data = argoverse_loader.get(log_id) city_name = argoverse_data.city_name avm = ArgoverseMap() sdb = SynchronizationDB(dataset_dir, collect_single_log_id=log_id) try: path = os.path.join(output_dir, log_id, 'road_gt') command = "rm -r " + path os.system(command) except BaseException: pass try: os.makedirs(os.path.join(output_dir, log_id, 'road_gt')) except BaseException: pass ply_fpath = os.path.join(dataset_dir, log_id, 'lidar') ply_locs = [] for idx, ply_name in enumerate(os.listdir(ply_fpath)): ply_loc = np.array([idx, int(ply_name.split('.')[0].split('_')[-1])]) ply_locs.append(ply_loc) ply_locs = np.array(ply_locs) lidar_timestamps = sorted(ply_locs[:, 1]) calib_path = f"{dataset_dir}/{log_id}/vehicle_calibration_info.json" calib_data = read_json_file(calib_path) ind = 0 for i in range(9): if calib_data['camera_data_'][i]['key'] ==\ 'image_raw_stereo_front_left': ind = i break rotation = np.array(calib_data['camera_data_'][ind]['value'] ['vehicle_SE3_camera_']['rotation']['coefficients']) translation = np.array(calib_data['camera_data_'][ind]['value'] ['vehicle_SE3_camera_']['translation']) egovehicle_SE3_cam = SE3(rotation=quat2rotmat(rotation), translation=translation) for idx in range(len(lidar_timestamps)): lidar_timestamp = lidar_timestamps[idx] cam_timestamp = sdb.get_closest_cam_channel_timestamp( lidar_timestamp, "stereo_front_left", str(log_id)) occupancy_map = np.zeros((256, 256)) pose_fpath = os.path.join( dataset_dir, log_id, "poses", "city_SE3_egovehicle_" + str(lidar_timestamp) + ".json") if not Path(pose_fpath).exists(): continue pose_data = read_json_file(pose_fpath) rotation = np.array(pose_data["rotation"]) translation = np.array(pose_data["translation"]) xcenter = translation[0] ycenter = translation[1] city_to_egovehicle_se3 = SE3(rotation=quat2rotmat(rotation), translation=translation) ego_car_nearby_lane_ids = avm.get_lane_ids_in_xy_bbox( xcenter, ycenter, city_name, 50.0) occupancy_map = get_lane_bev(ego_car_nearby_lane_ids, avm, city_name, city_to_egovehicle_se3, egovehicle_SE3_cam, occupancy_map, res, 255) output_loc = os.path.join( output_dir, log_id, 'road_gt', 'stereo_front_left_' + str(cam_timestamp) + '.png') cv2.imwrite(output_loc, occupancy_map)
def draw_ground_pts_in_image( sdb: SynchronizationDB, lidar_points: np.ndarray, city_to_egovehicle_se3: SE3, dataset_map: ArgoverseMap, log_id: str, lidar_timestamp: int, city_name: str, dataset_dir: str, experiment_prefix: str, plot_ground: bool = True, motion_compensate: bool = False, camera: Optional[str] = None, ) -> Union[None, np.ndarray]: """Write an image to disk with rendered ground points for every camera. Args: sdb: instance of SynchronizationDB lidar_points: Numpy array of shape (N,3) in egovehicle frame city_to_egovehicle_se3: SE3 instance which takes a point in egovehicle frame and brings it into city frame dataset_map: Map dataset instance log_id: ID of the log city_name: A city's name (e.g. 'MIA' or 'PIT') motion_compensate: Whether to bring lidar points from world frame @ lidar timestamp, to world frame @ camera timestamp camera: camera name, if specified will return image of that specific camera, if None, will save all camera to disk and return None """ # put into city coords, then prune away ground and non-RoI points lidar_points = city_to_egovehicle_se3.transform_point_cloud(lidar_points) lidar_points = dataset_map.remove_non_driveable_area_points( lidar_points, city_name) _, not_ground_logicals = dataset_map.remove_ground_surface( copy.deepcopy(lidar_points), city_name, return_logicals=True) lidar_points = lidar_points[np.logical_not(not_ground_logicals) if plot_ground else not_ground_logicals] # put back into ego-vehicle coords lidar_points = city_to_egovehicle_se3.inverse_transform_point_cloud( lidar_points) calib_fpath = f"{dataset_dir}/{log_id}/vehicle_calibration_info.json" calib_data = read_json_file(calib_fpath) # repeat green to red colormap every 50 m. colors_arr = np.array([ [color_obj.rgb] for color_obj in Color("red").range_to(Color("green"), NUM_RANGE_BINS) ]).squeeze() np.fliplr(colors_arr) for cam_idx, camera_name in enumerate(RING_CAMERA_LIST + STEREO_CAMERA_LIST): im_dir = f"{dataset_dir}/{log_id}/{camera_name}" # load images, e.g. 'image_raw_ring_front_center_000000486.jpg' cam_timestamp = sdb.get_closest_cam_channel_timestamp( lidar_timestamp, camera_name, log_id) if cam_timestamp is None: continue im_fname = f"{camera_name}_{cam_timestamp}.jpg" im_fpath = f"{im_dir}/{im_fname}" # 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() points_h = point_cloud_to_homogeneous(copy.deepcopy(lidar_points)).T if motion_compensate: uv, uv_cam, valid_pts_bool = project_lidar_to_img_motion_compensated( points_h, # these are recorded at lidar_time copy.deepcopy(calib_data), camera_name, cam_timestamp, lidar_timestamp, dataset_dir, log_id, False, ) else: uv, uv_cam, valid_pts_bool = project_lidar_to_img( points_h, copy.deepcopy(calib_data), camera_name, False) if valid_pts_bool is None or uv is None or uv_cam is None: continue if valid_pts_bool.sum() == 0: continue uv = np.round(uv[valid_pts_bool]).astype(np.int32) uv_cam = uv_cam.T[valid_pts_bool] pt_ranges = np.linalg.norm(uv_cam[:, :3], axis=1) rgb_bins = np.round(pt_ranges).astype(np.int32) # account for moving past 100 meters, loop around again rgb_bins = rgb_bins % NUM_RANGE_BINS uv_colors = (255 * colors_arr[rgb_bins]).astype(np.int32) img = draw_point_cloud_in_img_cv2(img, uv, np.fliplr(uv_colors)) if not Path(f"{experiment_prefix}_ground_viz/{log_id}/{camera_name}" ).exists(): os.makedirs( f"{experiment_prefix}_ground_viz/{log_id}/{camera_name}") save_dir = f"{experiment_prefix}_ground_viz/{log_id}/{camera_name}" cv2.imwrite(f"{save_dir}/{camera_name}_{lidar_timestamp}.jpg", img) if camera == camera_name: return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) return None
class PerFrameLabelAccumulator: """We will cache the accumulated track label trajectories per city, per log, and per frame. In order to plot each frame sequentially, one at a time, we need to aggregate beforehand the tracks and cuboids for each frame. Attributes: bboxes_3d (bool): to use 3d bounding boxes (True) or 2d bounding boxes (False). dataset_dir (str): Dataset directory. labels_dir (str): Labels directory. log_egopose_dict (dict): Egopose per log id and timestamp. log_timestamp_dict (dict): List of frame records per log id and timestamp. per_city_traj_dict (dict): Per city trajectory dictionary. sdb (SynchronizationDB): Synchronization DB. """ def __init__( self, dataset_dir: str, labels_dir: str, experiment_prefix: str, bboxes_3d: bool = False, save: bool = True, ) -> None: """Initialize PerFrameLabelAccumulator object for use with tracking benchmark data. Args: dataset_dir (str): Dataset directory. labels_dir (str): Labels directory. experiment_prefix (str): Prefix for experimint to use. bboxes_3d (bool, optional): to use 3d bounding boxes (True) or 2d bounding boxes (False). """ self.bboxes_3d = bboxes_3d self.dataset_dir = dataset_dir self.labels_dir = labels_dir tmp_dir = tempfile.gettempdir() per_city_traj_dict_fpath = f"{tmp_dir}/per_city_traj_dict_{experiment_prefix}.pkl" log_egopose_dict_fpath = f"{tmp_dir}/log_egopose_dict_{experiment_prefix}.pkl" log_timestamp_dict_fpath = f"{tmp_dir}/log_timestamp_dict_{experiment_prefix}.pkl" # coordinate system is the map world frame self.per_city_traj_dict: Dict[str, List[Tuple[np.ndarray, str]]] = { "MIA": [], "PIT": [], } # all the trajectories for these 2 cities self.log_egopose_dict: Dict[str, Dict[int, Dict[str, np.ndarray]]] = {} self.log_timestamp_dict: Dict[str, Dict[int, List[FrameRecord]]] = {} self.sdb = SynchronizationDB(self.dataset_dir) if save: self.accumulate_per_log_data() save_pkl_dictionary(per_city_traj_dict_fpath, self.per_city_traj_dict) save_pkl_dictionary(log_egopose_dict_fpath, self.log_egopose_dict) save_pkl_dictionary(log_timestamp_dict_fpath, self.log_timestamp_dict) def accumulate_per_log_data(self, log_id: Optional[str] = None) -> None: """Loop through all of the logs that we have. Get the labels that pertain to the benchmark (i.e. tracking or detection) that we are interested in. We use a unique color to describe each trajectory, and then we store the instance of the trajectory, along with its color, *PER FRAME* , per log. """ MIAMI_CUBOID_COUNT = 0 PITT_CUBOID_COUNT = 0 log_fpaths = glob.glob(f"{self.dataset_dir}/*") log_fpaths = [f for f in log_fpaths if os.path.isdir(f)] num_benchmark_logs = len(log_fpaths) for log_idx, log_fpath in enumerate(log_fpaths): log_id_ = log_fpath.split("/")[-1] if log_id is not None: if log_id_ != log_id: continue if log_id_ not in self.sdb.get_valid_logs(): continue city_info_fpath = f"{self.dataset_dir}/{log_id_}/city_info.json" city_info = read_json_file(city_info_fpath) log_city_name = city_info["city_name"] if log_city_name not in self.per_city_traj_dict: logger.warning(f"{log_city_name} not listed city") continue self.log_egopose_dict[log_id_] = {} self.log_timestamp_dict[log_id_] = {} traj_labels = self.get_log_trajectory_labels(log_id_) if traj_labels is None: continue # skip this log since no tracking data for traj_idx, traj_label in enumerate(traj_labels): if (traj_idx % 500) == 0: logger.info(f"On traj index {traj_idx}") traj_city_fr = self.place_trajectory_in_city_frame( traj_label, log_id_) # we don't know the city name until here if traj_idx == 0: logger.info( f"Log {log_id_} has {len(traj_labels)} trajectories in {log_city_name}" ) self.per_city_traj_dict[log_city_name].append( (traj_city_fr, log_id_)) logger.info(f"We looked at {num_benchmark_logs} tracking logs") logger.info( f"Miami has {MIAMI_CUBOID_COUNT} and Pittsburgh has {PITT_CUBOID_COUNT} cuboids" ) def get_log_trajectory_labels( self, log_id: str) -> Optional[List[TrajectoryLabel]]: """Create a very large list with all of the trajectory data. Treat a single object cuboid label as one step in a trajectory. Then we can share the same representation for both. Args: log_id (str): Log id to load. Returns: List[TrajectoryLabel]: List of trajectory labels. """ path = f"{self.labels_dir}/{log_id}/track_labels_amodal" if Path(path).exists(): return load_json_track_labels(f"{path}/*.json") else: return None def place_trajectory_in_city_frame(self, traj_label: TrajectoryLabel, log_id: str) -> np.ndarray: """Place trajectory in the city frame Args: traj_label (TrajectoryLabel): instance of the TrajectoryLabel class. log_id (str): Log id. Returns: - traj_city_fr: trajectory length of NUM_CUBOID_VERTS (x,y,z) coords per cuboid. """ seq_len = traj_label.timestamps.shape[0] if self.bboxes_3d: NUM_CUBOID_VERTS = 8 else: NUM_CUBOID_VERTS = 4 # store NUM_CUBOID_VERTS (x,y,z) coords per cuboid traj_city_fr = np.zeros((seq_len, NUM_CUBOID_VERTS, 3)) rand_color = ( float(np.random.rand()), float(np.random.rand()), float(np.random.rand()), ) logger.info(f"On log {log_id} with {traj_label.track_uuid}") for t in range(seq_len): obj_label_rec = ObjectLabelRecord( quaternion=traj_label.quaternions[t], translation=traj_label.translations[t], length=traj_label.max_length, width=traj_label.max_width, height=traj_label.max_height, occlusion=traj_label.occlusion[t], ) timestamp = int(traj_label.timestamps[t]) if self.bboxes_3d: bbox_ego_frame = obj_label_rec.as_3d_bbox() else: bbox_ego_frame = obj_label_rec.as_2d_bbox() bbox_city_fr, pose_city_to_ego = self.convert_bbox_to_city_frame( timestamp, self.dataset_dir, log_id, bbox_ego_frame) if bbox_city_fr is None: logger.warning( f"\t {log_id}: Couldnt find the pose for {traj_label.track_uuid}!" ) continue self.log_egopose_dict[log_id][timestamp] = pose_city_to_ego frame_rec = FrameRecord( bbox_city_fr=bbox_city_fr, bbox_ego_frame=bbox_ego_frame, occlusion_val=obj_label_rec.occlusion, color=rand_color, track_uuid=traj_label.track_uuid, obj_class_str=traj_label.obj_class_str, ) self.log_timestamp_dict[log_id].setdefault(timestamp, []).append(frame_rec) traj_city_fr[t] = bbox_city_fr return traj_city_fr def convert_bbox_to_city_frame( self, lidar_timestamp_ns: int, dataset_dir: str, log_id: str, bbox_ego_frame: np.ndarray, ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: """Convert bounding box to city frame. Args: lidar_timestamp_ns (int): Lidar timestamp. dataset_dir (str): representing full path to the log_ids. log_id (str): e.g. '3ced8dba-62d0-3930-8f60-ebeea2feabb8'. bbox_ego_frame (np.ndarray): Numpy array of shape (4,3), representing bounding box in egovehicle frame Returned: bbox_city_fr: Numpy array of shape (4,3), representing bounding box in CITY frame pose_city_to_ego: dictionary, has two fields: 'translation' and 'rotation' describing the SE(3) for p_city = city_to_egovehicle_se3 * p_egovehicle """ city_to_egovehicle_se3 = get_city_SE3_egovehicle_at_sensor_t( lidar_timestamp_ns, dataset_dir, log_id) if city_to_egovehicle_se3 is None: raise RuntimeError( f"Could not get city to egovehicle coordinate transformation at timestamp {lidar_timestamp_ns}" ) bbox_city_fr = city_to_egovehicle_se3.transform_point_cloud( bbox_ego_frame) pose_city_to_ego = { "rotation": city_to_egovehicle_se3.rotation, "translation": city_to_egovehicle_se3.translation, } return bbox_city_fr, pose_city_to_ego
class SimpleArgoverseTrackingDataLoader: """ Simple abstraction for retrieving log data, given a path to the dataset. """ def __init__(self, data_dir: str, labels_dir: str) -> None: """ Args: data_dir: str, representing path to raw Argoverse data labels_dir: strrepresenting path to Argoverse data labels """ self.data_dir = data_dir self.labels_dir = labels_dir self.sdb = SynchronizationDB(data_dir) def get_city_name(self, log_id: str) -> str: """ Args: log_id: str Returns: city_name: str """ city_info_fpath = f"{self.data_dir}/{log_id}/city_info.json" city_info = read_json_file(city_info_fpath) city_name = city_info["city_name"] assert isinstance(city_name, str) return city_name def get_log_calibration_data(self, log_id: str) -> Mapping[str, Any]: """ Args: log_id: str Returns: log_calib_data: dictionary """ calib_fpath = f"{self.data_dir}/{log_id}/vehicle_calibration_info.json" log_calib_data = read_json_file(calib_fpath) assert isinstance(log_calib_data, dict) return log_calib_data def get_city_to_egovehicle_se3(self, log_id: str, timestamp: int) -> Optional[SE3]: """ Args: log_id: str, unique ID of vehicle log timestamp: int, timestamp of sensor observation, in nanoseconds Returns: city_to_egovehicle_se3: SE3 transformation to bring egovehicle frame point into city frame. """ pose_fpath = f"{self.data_dir}/{log_id}/poses/city_SE3_egovehicle_{timestamp}.json" if not Path(pose_fpath).exists(): return None pose_data = read_json_file(pose_fpath) rotation = np.array(pose_data["rotation"]) translation = np.array(pose_data["translation"]) city_to_egovehicle_se3 = SE3(rotation=quat2rotmat(rotation), translation=translation) return city_to_egovehicle_se3 def get_closest_im_fpath(self, log_id: str, camera_name: str, lidar_timestamp: int) -> Optional[str]: """ Args: log_id: str, unique ID of vehicle log camera_name: str lidar_timestamp: int, timestamp of LiDAR sweep capture, in nanoseconds Returns: im_fpath, string representing path to image, or else None. """ cam_timestamp = self.sdb.get_closest_cam_channel_timestamp( lidar_timestamp, camera_name, log_id) if cam_timestamp is None: return None im_dir = f"{self.data_dir}/{log_id}/{camera_name}" im_fname = f"{camera_name}_{cam_timestamp}.jpg" im_fpath = f"{im_dir}/{im_fname}" return im_fpath def get_closest_lidar_fpath(self, log_id: str, cam_timestamp: int) -> Optional[str]: """ Args: log_id: str, unique ID of vehicle log cam_timestamp: int, timestamp of image capture, in nanoseconds Returns: ply_fpath: str, string representing path to PLY file, or else None. """ lidar_timestamp = self.sdb.get_closest_lidar_timestamp( cam_timestamp, log_id) if lidar_timestamp is None: return None lidar_dir = f"{self.data_dir}/{log_id}/lidar" ply_fname = f"PC_{lidar_timestamp}.ply" ply_fpath = f"{lidar_dir}/{ply_fname}" return ply_fpath def get_ordered_log_ply_fpaths(self, log_id: str) -> List[str]: """ Args: log_id: str, unique ID of vehicle log Returns: ply_fpaths: List of strings, representing paths to ply files in this log """ ply_fpaths = sorted( glob.glob(f"{self.data_dir}/{log_id}/lidar/PC_*.ply")) return ply_fpaths def get_ordered_log_cam_fpaths(self, log_id: str, camera_name: str) -> List[str]: """ Args log_id: str, unique ID of vehicle log Returns cam_img_fpaths: List of strings, representing paths to JPEG files in this log, for a specific camera """ cam_img_fpaths = sorted( glob.glob( f"{self.data_dir}/{log_id}/{camera_name}/{camera_name}_*.jpg")) return cam_img_fpaths def get_labels_at_lidar_timestamp( self, log_id: str, lidar_timestamp: int) -> Optional[List[Mapping[str, Any]]]: """ Args: log_id: str, unique ID of vehicle log lidar_timestamp: int, timestamp of LiDAR sweep capture, in nanoseconds Returns: labels: dictionary """ timestamp_track_label_fpath = ( f"{self.labels_dir}/{log_id}/per_sweep_annotations_amodal/tracked_object_labels_{lidar_timestamp}.json" ) if not Path(timestamp_track_label_fpath).exists(): return None labels = read_json_file(timestamp_track_label_fpath) assert isinstance(labels, list), labels return labels
class SimpleArgoverseTrackingDataLoader: """ Simple abstraction for retrieving log data, given a path to the dataset. """ def __init__(self, data_dir: str, labels_dir: str) -> None: """ Args: data_dir: str, representing path to raw Argoverse data labels_dir: str representing path to Argoverse data labels (e.g. labels or estimated detections/tracks) """ self.data_dir = data_dir self.labels_dir = labels_dir self.sdb = SynchronizationDB(data_dir) def get_city_name(self, log_id: str) -> str: """ Args: log_id: str Returns: city_name: str """ city_info_fpath = f"{self.data_dir}/{log_id}/city_info.json" city_name = read_city_name(city_info_fpath) assert isinstance(city_name, str) return city_name def get_log_calibration_data(self, log_id: str) -> Mapping[str, Any]: """ Args: log_id: str Returns: log_calib_data: dictionary """ calib_fpath = f"{self.data_dir}/{log_id}/vehicle_calibration_info.json" log_calib_data = read_json_file(calib_fpath) assert isinstance(log_calib_data, dict) return log_calib_data def get_city_to_egovehicle_se3(self, log_id: str, timestamp: int) -> Optional[SE3]: """Deprecated version of get_city_SE3_egovehicle() below, as does not follow standard naming convention Args: log_id: str, unique ID of vehicle log timestamp: int, timestamp of sensor observation, in nanoseconds Returns: city_SE3_egovehicle: SE3 transformation to bring points in egovehicle frame into city frame. """ return self.get_city_SE3_egovehicle(log_id, timestamp) def get_city_SE3_egovehicle(self, log_id: str, timestamp: int) -> Optional[SE3]: """ Args: log_id: str, unique ID of vehicle log timestamp: int, timestamp of sensor observation, in nanoseconds Returns: city_SE3_egovehicle: SE3 transformation to bring points in egovehicle frame into city frame. """ return get_city_SE3_egovehicle_at_sensor_t(timestamp, self.data_dir, log_id) def get_closest_im_fpath(self, log_id: str, camera_name: str, lidar_timestamp: int) -> Optional[str]: """ Args: log_id: str, unique ID of vehicle log camera_name: str lidar_timestamp: int, timestamp of LiDAR sweep capture, in nanoseconds Returns: im_fpath, string representing path to image, or else None. """ cam_timestamp = self.sdb.get_closest_cam_channel_timestamp( lidar_timestamp, camera_name, log_id) if cam_timestamp is None: return None im_dir = f"{self.data_dir}/{log_id}/{camera_name}" im_fname = f"{camera_name}_{cam_timestamp}.jpg" im_fpath = f"{im_dir}/{im_fname}" return im_fpath def get_closest_lidar_fpath(self, log_id: str, cam_timestamp: int) -> Optional[str]: """ Args: log_id: str, unique ID of vehicle log cam_timestamp: int, timestamp of image capture, in nanoseconds Returns: ply_fpath: str, string representing path to PLY file, or else None. """ lidar_timestamp = self.sdb.get_closest_lidar_timestamp( cam_timestamp, log_id) if lidar_timestamp is None: return None lidar_dir = f"{self.data_dir}/{log_id}/lidar" ply_fname = f"PC_{lidar_timestamp}.ply" ply_fpath = f"{lidar_dir}/{ply_fname}" return ply_fpath def get_ordered_log_ply_fpaths(self, log_id: str) -> List[str]: """ Args: log_id: str, unique ID of vehicle log Returns: ply_fpaths: List of strings, representing paths to chronologically ordered ply files in this log File paths are strings are of the same length ending with a nanosecond timestamp, thus sorted() will place them in numerical order. """ ply_fpaths = sorted( glob.glob(f"{self.data_dir}/{log_id}/lidar/PC_*.ply")) return ply_fpaths def get_ordered_log_cam_fpaths(self, log_id: str, camera_name: str) -> List[str]: """ Args log_id: str, unique ID of vehicle log Returns cam_img_fpaths: List of strings, representing paths to ordered JPEG files in this log, for a specific camera """ cam_img_fpaths = sorted( glob.glob( f"{self.data_dir}/{log_id}/{camera_name}/{camera_name}_*.jpg")) return cam_img_fpaths def get_labels_at_lidar_timestamp( self, log_id: str, lidar_timestamp: int) -> Optional[List[Mapping[str, Any]]]: """ Args: log_id: str, unique ID of vehicle log lidar_timestamp: int, timestamp of LiDAR sweep capture, in nanoseconds Returns: labels: dictionary """ timestamp_track_label_fpath = ( f"{self.labels_dir}/{log_id}/per_sweep_annotations_amodal/tracked_object_labels_{lidar_timestamp}.json" ) if not Path(timestamp_track_label_fpath).exists(): return None labels = read_json_file(timestamp_track_label_fpath) assert isinstance(labels, list), labels return labels