def dump_object_labels( labels: List[waymo_open_dataset.label_pb2.Label], timestamp: int, log_id: str, parent_path: str, track_id_dict: Dict, ) -> None: """Saves object labels from Waymo dataset as json files Args: labels: A list of Waymo labels timestamp: Timestamp in nanoseconds when the lidar reading occurred log_id: Log ID that the reading belongs to parent_path: The directory that the converted data is written to track_id_dict: Dictionary to store object ID to track ID mappings """ argoverse_labels = [] for label in labels: # We don't want signs, as that is not a category in Argoverse if label.type != LABEL_TYPES.index("SIGN") and label.type != LABEL_TYPES.index( "UNKNOWN" ): argoverse_labels.append(build_argo_label(label, timestamp, track_id_dict)) json_fpath = f"{parent_path}/{log_id}/per_sweep_annotations_amodal/" json_fpath += f"tracked_object_labels_{timestamp}.json" check_mkdir(str(Path(json_fpath).parent)) save_json_dict(json_fpath, argoverse_labels)
def save_to_disk(self) -> None: """ Labels and predictions should be saved in JSON e.g. `tracked_object_labels_315969629019741000.json` """ for ts_ns, ts_trackedlabels in self.ts_to_trackedlabels_dict.items(): json_fpath = f"{self.log_dir}/per_sweep_annotations_amodal/" check_mkdir(json_fpath) json_fpath += f"tracked_object_labels_{ts_ns}.json" save_json_dict(json_fpath, ts_trackedlabels)
def save_as_json(self, save_fpath: _PathLike) -> None: """Save the Sim(2) object to a JSON representation on disk. Args: save_fpath: path to where json file should be saved """ dict_for_serialization = { "R": self.rotation.flatten().tolist(), "t": self.translation.flatten().tolist(), "s": self.scale, } save_json_dict(save_fpath, dict_for_serialization)
def test_save_json_dict() -> None: """Test saving a dictionary to a JSON file.""" json_fpath = _TEST_DIR / "test_data/json_save_test_file.json" intended_dict = {"a": 1, "b": None, "c": 9999, "d": "abc"} save_json_dict(json_fpath, intended_dict) # verify it was saved correctly with open(json_fpath, "rb") as f: loaded_dict = json.load(f) assert intended_dict == loaded_dict # make sure key sets are identical assert set(intended_dict.keys()) == set(loaded_dict.keys()) for k in intended_dict.keys(): assert intended_dict[k] == loaded_dict[k]
def dump_pose( city_SE3_egovehicle: np.ndarray, timestamp: int, log_id: str, parent_path: str ) -> None: """Saves the SE3 transformation from city frame to egovehicle frame at a particular timestamp Args: city_SE3_egovehicle: A (4,4) numpy array representing the SE3 transformation from city to egovehicle frame timestamp: Timestamp in nanoseconds when the lidar reading occurred log_id: Log ID that the reading belongs to parent_path: The directory that the converted data is written to """ x, y, z = city_SE3_egovehicle[:3, 3] R = city_SE3_egovehicle[:3, :3] assert np.allclose(city_SE3_egovehicle[3], np.array([0, 0, 0, 1])) q = rotmat2quat(R) w, x, y, z = q pose_dict = {"rotation": [w, x, y, z], "translation": [x, y, z]} json_fpath = f"{parent_path}/{log_id}/poses/city_SE3_egovehicle_{timestamp}.json" check_mkdir(str(Path(json_fpath).parent)) save_json_dict(json_fpath, pose_dict)
list_path_label_persweep = glob.glob( os.path.join(path_log, "per_sweep_annotations_amodal", "*")) list_path_label_persweep.sort() dist_track_labels = {} for path_label_persweep in list_path_label_persweep: data = read_json_file(path_label_persweep) for data_obj in data: id_obj = data_obj['track_label_uuid'] if id_obj not in dist_track_labels.keys(): dist_track_labels[id_obj] = [] dist_track_labels[id_obj].append(data_obj) path_amodal_labels = os.path.join(path_log, "track_labels_amodal") data_amodal = {} if os.path.exists(path_amodal_labels): shutil.rmtree(path_amodal_labels) os.mkdir(path_amodal_labels) print("Adding files to ", path_amodal_labels) for key in dist_track_labels.keys(): data_amodal[key] = {} data_amodal[key]['label_class'] = dist_track_labels[key][0][ 'label_class'] data_amodal[key]['uuid'] = dist_track_labels[key][0][ 'track_label_uuid'] data_amodal[key]['log_id'] = path_log.split('/')[-1] data_amodal[key]['track_label_frames'] = dist_track_labels[key] save_json_dict(os.path.join(path_amodal_labels, "%s.json" % key), data_amodal[key])
def main(args: argparse.Namespace) -> None: """Main script to convert Waymo object labels, LiDAR, images, pose, and calibration to the Argoverse data format on disk """ TFRECORD_DIR = args.waymo_dir ARGO_WRITE_DIR = args.argo_dir track_id_dict = {} img_count = 0 log_ids = get_log_ids_from_files(TFRECORD_DIR) for log_id, tf_fpath in log_ids.items(): dataset = tf.data.TFRecordDataset(tf_fpath, compression_type="") log_calib_json = None for data in dataset: frame = open_dataset.Frame() frame.ParseFromString(bytearray(data.numpy())) # Checking if we extracted the correct log ID assert log_id == frame.context.name # Frame start time, which is the timestamp # of the first top lidar spin within this frame, in microseconds timestamp_ms = frame.timestamp_micros timestamp_ns = int(timestamp_ms * 1000) # to nanoseconds SE3_flattened = np.array(frame.pose.transform) city_SE3_egovehicle = SE3_flattened.reshape(4, 4) if args.save_poses: dump_pose(city_SE3_egovehicle, timestamp_ns, log_id, ARGO_WRITE_DIR) # Reading lidar data and saving it in point cloud format # We are only using the first range image (Waymo provides two range images) # If you want to use the second one, you can change it in the arguments ( range_images, camera_projections, range_image_top_pose, ) = frame_utils.parse_range_image_and_camera_projection(frame) if args.range_image == 1: points_ri, cp_points_ri = frame_utils.convert_range_image_to_point_cloud( frame, range_images, camera_projections, range_image_top_pose ) elif args.range_image == 2: points_ri, cp_points_ri = frame_utils.convert_range_image_to_point_cloud( frame, range_images, camera_projections, range_image_top_pose, ri_index=1, ) points_all_ri = np.concatenate(points_ri, axis=0) if args.save_cloud: dump_point_cloud(points_all_ri, timestamp_ns, log_id, ARGO_WRITE_DIR) # Saving labels if args.save_labels: dump_object_labels( frame.laser_labels, timestamp_ns, log_id, ARGO_WRITE_DIR, track_id_dict, ) if args.save_calibration: calib_json = form_calibration_json(frame.context.camera_calibrations) if log_calib_json is None: log_calib_json = calib_json calib_json_fpath = ( f"{ARGO_WRITE_DIR}/{log_id}/vehicle_calibration_info.json" ) check_mkdir(str(Path(calib_json_fpath).parent)) save_json_dict(calib_json_fpath, calib_json) else: assert calib_json == log_calib_json # 5 images per frame for index, tf_cam_image in enumerate(frame.images): # 4x4 row major transform matrix that transforms # 3d points from one frame to another. SE3_flattened = np.array(tf_cam_image.pose.transform) city_SE3_egovehicle = SE3_flattened.reshape(4, 4) # in seconds timestamp_s = tf_cam_image.pose_timestamp # TODO: this looks buggy, need to confirm timestamp_ns = int(round_to_micros(int(timestamp_s * 1e9)) * 1000) # to nanoseconds if args.save_poses: dump_pose(city_SE3_egovehicle, timestamp_ns, log_id, ARGO_WRITE_DIR) if args.save_images: camera_name = CAMERA_NAMES[tf_cam_image.name] img = tf.image.decode_jpeg(tf_cam_image.image) new_img = undistort_image( np.asarray(img), frame.context.camera_calibrations, tf_cam_image.name, ) img_save_fpath = f"{ARGO_WRITE_DIR}/{log_id}/{camera_name}/" img_save_fpath += f"{camera_name}_{timestamp_ns}.jpg" check_mkdir(str(Path(img_save_fpath).parent)) imageio.imwrite(img_save_fpath, new_img) img_count += 1 if img_count % 100 == 0: print(f"\tSaved {img_count}'th image for log = {log_id}")
def main(nusc: NuScenes, args: argparse.Namespace, start_index: int, end_index: int) -> None: """ Convert sweeps and samples into (unannotated) Argoverse format. Overview of algorithm: 1) Iterate over all scenes in the NuScenes dataset. For each scene, obtain first sample in the scene. 2) Get the sample_data corresponding to each of the channels from the sample, and convert it to argo format. 3) While the sample_data is not corresponding to a key_frame, get the next sample_data, and repeat step 2. 4) Go to the next sample while we are in the same scene. """ OUTPUT_ROOT = args.argo_dir NUSCENES_ROOT = args.nuscenes_dir NUSCENES_VERSION = args.nuscenes_version if not os.path.exists(OUTPUT_ROOT): os.makedirs(OUTPUT_ROOT) tot_scenes = len(nusc.scene) for scene in nusc.scene[start_index:min(end_index, tot_scenes)]: scene_token = scene["token"] sample_token = scene["first_sample_token"] scene_path = os.path.join(OUTPUT_ROOT, scene_token) if not os.path.exists(scene_path): os.makedirs(scene_path) log_token = scene["log_token"] nusc_log = nusc.get("log", log_token) nusc_city = nusc_log["location"] save_json_dict(os.path.join(scene_path, f"city_info.json"), {"city_name": CITY_TO_ID[nusc_city]}) # Calibration info for all the sensors calibration_info = get_calibration_info(nusc, scene) calib_path = os.path.join(scene_path, f"vehicle_calibration_info.json") save_json_dict(calib_path, calibration_info) while sample_token != "": sample = nusc.get("sample", sample_token) timestamp = round_to_micros(sample["timestamp"]) tracked_labels = [] # city_SE3_vehicle pose ego_pose = None nsweeps_lidar = 10 nsweeps_cam = 6 # Save ego pose to json file poses_path = os.path.join(scene_path, f"poses") if not os.path.exists(poses_path): os.makedirs(poses_path) # Copy nuscenes sensor data into argoverse format and get the pose of the vehicle in the city frame for sensor, sensor_token in sample["data"].items(): if sensor in SENSOR_NAMES: argo_sensor = SENSOR_NAMES[sensor] output_sensor_path = os.path.join(scene_path, argo_sensor) if not os.path.exists(output_sensor_path): os.makedirs(output_sensor_path) sensor_data = nusc.get("sample_data", sensor_token) file_path = os.path.join(NUSCENES_ROOT, sensor_data["filename"]) i = 0 if sensor == "LIDAR_TOP": # nuscenes lidar data is stored as (x, y, z, intensity, ring index) while i < nsweeps_lidar and sensor_token != "": sensor_data = nusc.get("sample_data", sensor_token) file_path = os.path.join(NUSCENES_ROOT, sensor_data["filename"]) timestamp = round_to_micros( sensor_data["timestamp"]) # Not always exactly 10 if (sensor_data["is_key_frame"] and i != 0) or sample_token == "": break scan = np.fromfile(file_path, dtype=np.float32) points = scan.reshape((-1, 5)) # Transform lidar points from point sensor frame to egovehicle frame calibration = nusc.get( "calibrated_sensor", sensor_data["calibrated_sensor_token"], ) egovehicle_R_lidar = quat2rotmat( calibration["rotation"]) egovehicle_t_lidar = np.array( calibration["translation"]) egovehicle_SE3_lidar = SE3( rotation=egovehicle_R_lidar, translation=egovehicle_t_lidar, ) points_egovehicle = egovehicle_SE3_lidar.transform_point_cloud( points[:, :3]) write_ply(points_egovehicle, points, output_sensor_path, timestamp) if not os.path.isfile( os.path.join( poses_path, f"city_SE3_egovehicle_{timestamp}.json" )): ego_pose = nusc.get( "ego_pose", sensor_data["ego_pose_token"]) ego_pose_dict = { "rotation": ego_pose["rotation"], "translation": ego_pose["translation"], } save_json_dict( os.path.join( poses_path, f"city_SE3_egovehicle_{timestamp}.json" ), ego_pose_dict) sensor_token = sensor_data["next"] else: while i < nsweeps_cam and sensor_token != "": sensor_data = nusc.get("sample_data", sensor_token) file_path = os.path.join(NUSCENES_ROOT, sensor_data["filename"]) timestamp = round_to_micros( sensor_data["timestamp"]) # Not always exactly 6 if sensor_data["is_key_frame"] and i != 0: break shutil.copy( file_path, os.path.join(output_sensor_path, f"{argo_sensor}_{timestamp}.jpg"), ) sensor_token = sensor_data["next"] if not os.path.isfile( os.path.join( poses_path, f"city_SE3_egovehicle_{timestamp}.json" )): ego_pose = nusc.get( "ego_pose", sensor_data["ego_pose_token"]) ego_pose_dict = { "rotation": ego_pose["rotation"], "translation": ego_pose["translation"], } save_json_dict( os.path.join( poses_path, f"city_SE3_egovehicle_{timestamp}.json" ), ego_pose_dict) sample_token = sample["next"]