Analysis of Dataset Preprocessing Code in BEVerse

Analysis of Dataset Preprocessing Code in BEVerse

It mainly analyzes the preprocessing part of the Nuscenes dataset in the BEVerse project.

1. Dataset generation script

Detailed reference:BEVerse/docs/data_preparation.md

• 1.1 Data structure

BEVerse
├── mmdet3d
├── tools
├── configs
├── projects
├── data
│   ├── nuscenes
│   │   ├── maps
│   │   ├── samples
│   │   ├── sweeps
│   │   ├── v1.0-test
|   |   ├── v1.0-trainval

• 1.2 data_info generation script

python tools/create_data.py nuscenes --root-path ./data/nuscenes --out-dir ./data/nuscenes --extra-tag nuscenes

2. In-depth look at the generated script

File location :BEVerse/tools/create_data.py

Core processing : In the main function, according to different data types of the dataset, such as 'kitti', 'nuscenes', etc., different data preprocessing is performed, and it is packaged into an info file of key-value pairs and stored. This article mainly focuses on the nuscenes_data_prep() function analysis.

Main code :

def nuscenes_data_prep(root_path,
                       info_prefix,
                       version,
                       dataset_name,
                       out_dir,
                       max_sweeps=10):
    """Prepare data related to nuScenes dataset.

    Related data consists of '.pkl' files recording basic infos,
    2D annotations and groundtruth database.

    Args:
        root_path (str): Path of dataset root.
        info_prefix (str): The prefix of info filenames.
        version (str): Dataset version.
        dataset_name (str): The dataset class name.
        out_dir (str): Output directory of the groundtruth database info.
        max_sweeps (int): Number of input consecutive frames. Default: 10
    """

    import os
    info_save_path = 'data/nuscenes_infos'
    os.makedirs(info_save_path, exist_ok=True)

    nuscenes_converter.create_nuscenes_infos(
        root_path, info_prefix, version=version, max_sweeps=max_sweeps, info_save_path=info_save_path)

    if version == 'v1.0-test':
        info_test_path = osp.join(
            info_save_path, f'{info_prefix}_infos_test.pkl')
        nuscenes_converter.export_2d_annotation(
            root_path, info_test_path, version=version)
        return

    info_train_path = osp.join(
        info_save_path, f'{info_prefix}_infos_train.pkl')
    info_val_path = osp.join(info_save_path, f'{info_prefix}_infos_val.pkl')

    nuscenes_converter.export_2d_annotation(
        root_path, info_train_path, version=version)
    nuscenes_converter.export_2d_annotation(
        root_path, info_val_path, version=version)

Nuscenes data preprocessing mainly includes establishing storage paths, establishing data information (create_nuscenes_infos), exporting 2D annotations (export_2d_annotation), etc.

3. Core functions

3.1 Create data information

1. Function prototype
file location:BEVerse/tools/data_converter/nuscenes_converter.py

def create_nuscenes_infos(root_path,
                          info_prefix,
                          version='v1.0-trainval',
                          max_sweeps=10,
                          info_save_path=None):
    """Create info file of nuscene dataset.

    Given the raw data, generate its related info file in pkl format.

    Args:
        root_path (str): Path of the data root.
        info_prefix (str): Prefix of the info file to be generated.
        version (str): Version of the data.
            Default: 'v1.0-trainval'
        max_sweeps (int): Max number of sweeps.
            Default: 10
    """
    ## 1. 利用Nuscenes工具导入scene场景
    from nuscenes.nuscenes import NuScenes
    nusc = NuScenes(version=version, dataroot=root_path, verbose=True)
    from nuscenes.utils import splits
    available_vers = ['v1.0-trainval', 'v1.0-test', 'v1.0-mini']
    assert version in available_vers
    if version == 'v1.0-trainval':
        train_scenes = splits.train
        val_scenes = splits.val
    elif version == 'v1.0-test':
        train_scenes = splits.test
        val_scenes = []
    elif version == 'v1.0-mini':
        train_scenes = splits.mini_train
        val_scenes = splits.mini_val
    else:
        raise ValueError('unknown')
    
    # 2. filter existing scenes.
    available_scenes = get_available_scenes(nusc)
    available_scene_names = [s['name'] for s in available_scenes]
    train_scenes = list(
        filter(lambda x: x in available_scene_names, train_scenes))
    val_scenes = list(filter(lambda x: x in available_scene_names, val_scenes))
    train_scenes = set([
        available_scenes[available_scene_names.index(s)]['token']
        for s in train_scenes
    ])
    val_scenes = set([
        available_scenes[available_scene_names.index(s)]['token']
        for s in val_scenes
    ])

    test = 'test' in version
    if test:
        print('test scene: {}'.format(len(train_scenes)))
    else:
        print('train scene: {}, val scene: {}'.format(
            len(train_scenes), len(val_scenes)))
    
    ## 3. 填充训练集验证集信息
    train_nusc_infos, val_nusc_infos = _fill_trainval_infos(
        nusc, train_scenes, val_scenes, test, max_sweeps=max_sweeps)

    info_save_path = info_save_path if info_save_path else root_path

    metadata = dict(version=version)
    if test:
        print('test sample: {}'.format(len(train_nusc_infos)))
        data = dict(infos=train_nusc_infos, metadata=metadata)
        info_path = osp.join(info_save_path,
                             '{}_infos_test.pkl'.format(info_prefix))
        mmcv.dump(data, info_path)
    else:
        print('train sample: {}, val sample: {}'.format(
            len(train_nusc_infos), len(val_nusc_infos)))
        # train infos
        data = dict(infos=train_nusc_infos, metadata=metadata)
        info_path = osp.join(info_save_path,
                             '{}_infos_train.pkl'.format(info_prefix))
        mmcv.dump(data, info_path)
        # val infos
        data['infos'] = val_nusc_infos
        info_val_path = osp.join(info_save_path,
                                 '{}_infos_val.pkl'.format(info_prefix))
        mmcv.dump(data, info_val_path)
        # trainval infos
        trainval_nusc_infos = train_nusc_infos + val_nusc_infos
        data['infos'] = trainval_nusc_infos
        info_trainval_path = osp.join(info_save_path,
                                      '{}_infos_trainval.pkl'.format(info_prefix))
        mmcv.dump(data, info_trainval_path)

2. Process analysis
1) Taking the mini dataset as an example, first use the built-in tools of Nuscenes to generate train_scenes (8 training scenes) and val_scenes (2 verification scenes).
2) Then filter the existing scenes, extract tokens from the existing training and verification scenes, recompose the set, and store them in train_scenesand respectively val_scenes.
3) Generate training set and verification set information (_fill_trainval_infos) from the original data, extract various data and attributes, and store them in and in the form of key-value train_nusc_infospairs val_nusc_infos.
4) Finally, pack train_nusc_infos/val_nusc_infos and metadata (version key-value pairs) together, put them in data, and store them as pkl files through mmcv.dump.

3.2 Export 2D annotations

1. Function prototype

def export_2d_annotation(root_path, info_path, version, mono3d=True):
    """Export 2d annotation from the info file and raw data.

    Args:
        root_path (str): Root path of the raw data.
        info_path (str): Path of the info file.
        version (str): Dataset version.
        mono3d (bool): Whether to export mono3d annotation. Default: True.
    """
    # get bbox annotations for camera
    camera_types = [
        'CAM_FRONT',
        'CAM_FRONT_RIGHT',
        'CAM_FRONT_LEFT',
        'CAM_BACK',
        'CAM_BACK_LEFT',
        'CAM_BACK_RIGHT',
    ]
    nusc_infos = mmcv.load(info_path)['infos']
    nusc = NuScenes(version=version, dataroot=root_path, verbose=True)
    # info_2d_list = []
    cat2Ids = [
        dict(id=nus_categories.index(cat_name), name=cat_name)
        for cat_name in nus_categories
    ]
    coco_ann_id = 0
    coco_2d_dict = dict(annotations=[], images=[], categories=cat2Ids)
    for info in mmcv.track_iter_progress(nusc_infos):
        for cam in camera_types:
            cam_info = info['cams'][cam]
            coco_infos = get_2d_boxes(
                nusc,
                cam_info['sample_data_token'],
                visibilities=['', '1', '2', '3', '4'],
                mono3d=mono3d)
            (height, width, _) = mmcv.imread(cam_info['data_path']).shape
            coco_2d_dict['images'].append(
                dict(
                    file_name=cam_info['data_path'].split('data/nuscenes/')
                    [-1],
                    id=cam_info['sample_data_token'],
                    token=info['token'],
                    cam2ego_rotation=cam_info['sensor2ego_rotation'],
                    cam2ego_translation=cam_info['sensor2ego_translation'],
                    ego2global_rotation=info['ego2global_rotation'],
                    ego2global_translation=info['ego2global_translation'],
                    cam_intrinsic=cam_info['cam_intrinsic'],
                    width=width,
                    height=height))
            for coco_info in coco_infos:
                if coco_info is None:
                    continue
                # add an empty key for coco format
                coco_info['segmentation'] = []
                coco_info['id'] = coco_ann_id
                coco_2d_dict['annotations'].append(coco_info)
                coco_ann_id += 1
    if mono3d:
        json_prefix = f'{info_path[:-4]}_mono3d'
    else:
        json_prefix = f'{info_path[:-4]}'
    mmcv.dump(coco_2d_dict, f'{json_prefix}.coco.json')

2. Process analysis
1) First get nusc_infos, which is the infos information stored in 3.1.
2) Arrange 10 categories into id index values
​​3) Traverse the infos information, and then traverse multiple cameras separately, and store attribute information in each camera, such as sample_data_token, various rotation matrices, and label information into coco_2d_dict.

4. Generated file and description of its contents

Reference article : 3D Object Detection NuScenes Dataset

For details, refer to docs/datasets/nuscenes_det.mdthe English document. The directory structure of data is as follows. Here we mainly focus on .pklfiles, which are used for methods involving point clouds. .jsonis a Coco-style image-based approach, such as image-based 2D and 3D detection.

├── data
│   ├── nuscenes
│   │   ├── maps
│   │   ├── samples
│   │   ├── sweeps
│   │   ├── v1.0-test
|   |   ├── v1.0-trainval
│   │   ├── nuscenes_database
│   │   ├── nuscenes_infos_train.pkl
│   │   ├── nuscenes_infos_val.pkl
│   │   ├── nuscenes_infos_test.pkl
│   │   ├── nuscenes_dbinfos_train.pkl
│   │   ├── nuscenes_infos_train_mono3d.coco.json
│   │   ├── nuscenes_infos_val_mono3d.coco.json
│   │   ├── nuscenes_infos_test_mono3d.coco.json

• nuscenes_database/xxxxx.bin: Point cloud data contained in each 3D bounding box of the training dataset.
• nuscenes_infos_train.pkl: training data set information, contains two key values: metadata 和 infos. metadata is the basic information of the dataset itself, such as {'version': 'v1.0-trainval'}, while infos contains detailed information.
  • info['lidar_path']: The file path of lidar point cloud data.
  • info['token']: Sample data token.
  • info['sweeps']: scan information. 有别于samples关键帧.
    • info['sweeps'][i]['data_path']: The data path of the i-th scan.
    • info['sweeps'][i]['type']: type of scan data, such as "lidar".
    • info['sweeps'][i]['sample_data_token']: Sweep sample data token.
    • info['sweeps'][i]['sensor2ego_translation']: translation from current sensor to ego vehicle (1x3 list).
    • info['sweeps'][i]['sensor2ego_rotation']: The rotation from the current sensor to the ego vehicle (1x4 list in quaternion format).
    • info['sweeps'][i]['ego2global_translation']: Translation from ego vehicle to global coordinates (1x3 list).
    • info['sweeps'][i]['ego2global_rotation']: Rotation from ego vehicle to global coordinates (1x4 list in quaternion format).
    • info['sweeps'][i]['timestamp']: The timestamp of the scanned data.
    • info['sweeps'][i]['sensor2lidar_translation']: Translation (1x3 list) from current sensor (used to collect scan data) to lidar.
    • info['sweeps'][i]['sensor2lidar_rotation']: Rotation from current sensor (used to collect scan data) to lidar (1x4 list in quaternion format).
  • info['cams']: Camera calibration information. It contains six key values ​​corresponding to each camera: 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_FRONT_LEFT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT'.
  • info['lidar2ego_translation']: Lidar to ego translation (1x3 list).
  • info['lidar2ego_rotation']: rotation from lidar to ego vehicle (1x4 list in quaternion format).
  • info['ego2global_translation']: Translation from ego to global coordinates (1x3 list).
  • info['ego2global_rotation']: Rotation from ego vehicle to global coordinates (1x4 list in quaternion format).
  • info['timestamp']: The timestamp of the sample data.
  • info['gt_boxes']: 3D bounding boxes with 7 degrees of freedom, an Nx7 array.
  • info['gt_names']: The category of the 3D bounding box, a 1xN array.
  • info['gt_velocity']: The velocity of the 3D bounding box (no vertical measurement due to inaccuracy), an Nx2 array.
  • info['num_lidar_pts']: The number of lidar points contained in each 3D bounding box.
  • info['num_radar_pts']: The number of radar points contained in each 3D bounding box.
  • info['valid_flag']: Whether each bounding box is valid. In general, we only consider 至少一个3D boxes containing lidar or mmWave radar points as valid boxes