| import argparse |
|
|
| import numpy as np |
|
|
| from waymo_preprocess import WaymoProcessor |
|
|
| if __name__ == "__main__": |
| """ |
| Waymo Dataset preprocessing script |
| =========================== |
| |
| This script facilitates the preprocessing of the Waymo dataset |
| |
| Usage: |
| ------ |
| python preprocess.py \ |
| --data_root <path_to_waymo_data> \ |
| --target_dir <output_directory> \ |
| [additional_arguments] |
| |
| Example: |
| -------- |
| python preprocess.py --data_root data/waymo/raw/ --target_dir data/waymo/processed --split training --workers 3 --scene_ids 700 754 114 |
| |
| Arguments: |
| ---------- |
| --data_root (str): |
| The root directory where the Waymo dataset is stored. This is a required argument. |
| |
| --split (str): |
| Specifies the name of the data split. Default is set to "training". |
| |
| --target_dir (str): |
| Designates the directory where the processed data will be saved. This is a mandatory argument. |
| |
| --workers (int): |
| The number of processing threads. Default is set to 4. |
| |
| --scene_ids (list[int]): |
| List of specific scene IDs for processing. Should be integers separated by spaces. |
| |
| --split_file (str): |
| If provided, indicates the path to a file located in `data/waymo_splits` that contains the desired scene IDs. |
| |
| --start_idx (int): |
| Used in conjunction with `num_scenes` to generate a list of scene IDs when neither `scene_ids` nor `split_file` are provided. |
| |
| --num_scenes (int): |
| The total number of scenes to be processed. |
| |
| --process_keys (list[str]): |
| Denotes the types of data components to be processed. Options include but aren't limited to "images", "lidar", "calib", "pose", etc. |
| |
| Notes: |
| ------ |
| The logic of the script ensures that if specific scene IDs (`scene_ids`) are provided, they are prioritized. |
| If a split file (`split_file`) is indicated, it is utilized next. |
| If neither is available, the script uses the `start_idx` and `num_scenes` parameters to determine the scene IDs. |
| """ |
| parser = argparse.ArgumentParser(description="Data converter arg parser") |
| parser.add_argument( |
| "--data_root", type=str, required=True, help="root path of waymo dataset" |
| ) |
| parser.add_argument("--split", type=str, default="training", help="split name") |
| parser.add_argument( |
| "--target_dir", |
| type=str, |
| required=True, |
| help="output directory of processed data", |
| ) |
| parser.add_argument( |
| "--workers", type=int, default=4, help="number of threads to be used" |
| ) |
| |
| parser.add_argument( |
| "--scene_ids", |
| default=None, |
| type=int, |
| nargs="+", |
| help="scene ids to be processed, a list of integers separated by space. Range: [0, 798] for training, [0, 202] for validation", |
| ) |
| parser.add_argument( |
| "--split_file", type=str, default=None, help="Split file in data/waymo_splits" |
| ) |
| parser.add_argument( |
| "--start_idx", |
| type=int, |
| default=0, |
| help="If no scene id or split_file is given, use start_idx and num_scenes to generate scene_ids_list", |
| ) |
| parser.add_argument( |
| "--num_scenes", |
| type=int, |
| default=200, |
| help="number of scenes to be processed", |
| ) |
| parser.add_argument( |
| "--process_keys", |
| nargs="+", |
| default=[ |
| "images", |
| "lidar", |
| "calib", |
| "pose", |
| "dynamic_masks", |
| ], |
| ) |
| args = parser.parse_args() |
| if args.scene_ids is not None: |
| scene_ids_list = args.scene_ids |
| elif args.split_file is not None: |
| |
| split_file = open(args.split_file, "r").readlines()[1:] |
| scene_ids_list = [int(line.strip().split(",")[0]) for line in split_file] |
| else: |
| scene_ids_list = np.arange(args.start_idx, args.start_idx + args.num_scenes) |
|
|
| waymo_processor = WaymoProcessor( |
| load_dir=args.data_root, |
| save_dir=args.target_dir, |
| prefix=args.split, |
| process_keys=args.process_keys, |
| process_id_list=scene_ids_list, |
| workers=args.workers, |
| ) |
| if args.scene_ids is not None and args.workers == 1: |
| for scene_id in args.scene_ids: |
| waymo_processor.convert_one(scene_id) |
| else: |
| waymo_processor.convert() |
|
|
