Function waymo_data_prep

tools/create_data.py:210–250

Prepare the info file for waymo dataset. Args: root_path (str): Path of dataset root. info_prefix (str): The prefix of info filenames. out_dir (str): Output directory of the generated info file. workers (int): Number of threads to be used. max_sweeps (int

(root_path, info_prefix, version, out_dir, workers, max_sweeps=5)

def waymo_data_prep(root_path, info_prefix, version, out_dir, workers, max_sweeps=5):
    """Prepare the info file for waymo dataset.

    Args:
        root_path (str): Path of dataset root.
        info_prefix (str): The prefix of info filenames.
        out_dir (str): Output directory of the generated info file.
        workers (int): Number of threads to be used.
        max_sweeps (int): Number of input consecutive frames. Default: 5 \
            Here we store pose information of these frames for later use.
    """
    from tools.data_converter import waymo_converter as waymo

    splits = ["training", "validation", "testing"]

    for i, split in enumerate(splits):
        load_dir = osp.join(root_path, "waymo_format", split)
        if split == "validation":
            save_dir = osp.join(out_dir, "kitti_format", "training")
        else:
            save_dir = osp.join(out_dir, "kitti_format", split)
        converter = waymo.Waymo2KITTI(
            load_dir,
            save_dir,
            prefix=str(i),
            workers=workers,
            test_mode=(split == "test"),
        )
        converter.convert()
    # Generate waymo infos
    out_dir = osp.join(out_dir, "kitti_format")
    kitti.create_waymo_info_file(out_dir, info_prefix, max_sweeps=max_sweeps)

    create_groundtruth_database(
        "WaymoDataset",
        out_dir,
        info_prefix,
        f"{out_dir}/{info_prefix}_infos_train.pkl",
        relative_path=False,
        with_mask=False,
    )


parser = argparse.ArgumentParser(description="Data converter arg parser")