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"]
