BEVerse 中数据集预处理代码浅析

BEVerse 中数据集预处理代码浅析

主要针对BEVerse工程中 Nuscenes数据集预处理部分，进行解析。

一. 数据集生成脚本

详细参考： BEVerse/docs/data_preparation.md

• 1.1 数据结构

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

• 1.2 data_info生成脚本

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

二. 深入查看生成脚本

文件位置：BEVerse/tools/create_data.py

核心处理：main函数中根据dataset的不同数据类型，如’kitti’,'nuscenes’等，进行不同的数据预处理，将其打包成键值对的info文件存储下来。此文主要针对nuscenes_data_prep()函数解析。

主要代码：

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数据预处理，主要包括建立存储路径、建立数据信息（create_nuscenes_infos）、导出2D标注（export_2d_annotation）等。

三. 核心函数

3.1 建立数据信息

1. 函数原型
文件位置：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. 流程浅析
1）此处以mini数据集为例，首先利用Nuscenes自带工具，生成train_scenes（训练场景8个），val_scenes（验证场景2个）。
2）然后过滤已存在场景，将已存在的训练和验证场景，提取token，重新组成set，分别存入train_scenes 和 val_scenes。
3）从原始数据中生成训练集和验证集信息（_fill_trainval_infos），将各种数据和属性提取，并以键值对的方式存储在train_nusc_infos和val_nusc_infos。
4）最后将train_nusc_infos/val_nusc_infos 和 metadata（版本的键值对）包装到一起，放入到data中，经过mmcv.dump存放成pkl文件。

3.2 导出2D标注

1. 函数原型

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. 流程浅析
1）首先获取到nusc_infos，即3.1 中存储的infos信息。
2）将10个种类，编排成id索引值
3）遍历infos信息，再分别遍历多个相机，将每个相机中属性信息，如sample_data_token、各种旋转矩阵、标注信息存储到coco_2d_dict中。

四. 生成文件及其内容说明

参考文章：3D 目标检测 NuScenes 数据集

详情参考docs/datasets/nuscenes_det.md英文文档，data的目录结构如下，这里主要关注.pkl文件，用于涉及点云的方法。.json是Coco-style的基于图像的方法，例如基于图像的 2D 和 3D 检测。

├── 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：训练数据集的每个 3D 包围框中包含的点云数据。
• nuscenes_infos_train.pkl：训练数据集信息，含有两个键值： metadata 和 infos。metadata 数据集本身的基本信息，例如 {‘version’: ‘v1.0-trainval’}，而 infos 包含详细信息。
  • info[‘lidar_path’]：激光雷达点云数据的文件路径。
  • info[‘token’]：样本数据标记。
  • info[‘sweeps’]：扫描信息。有别于samples关键帧。
    • info[‘sweeps’][i][‘data_path’]：第 i 次扫描的数据路径。
    • info[‘sweeps’][i][‘type’]：扫描数据类型，例如“激光雷达”。
    • info[‘sweeps’][i][‘sample_data_token’]：扫描样本数据标记。
    • info[‘sweeps’][i][‘sensor2ego_translation’]：从当前传感器到自车的平移（1x3 列表）。
    • info[‘sweeps’][i][‘sensor2ego_rotation’]：从当前传感器到自车的旋转（四元数格式的 1x4 列表）。
    • info[‘sweeps’][i][‘ego2global_translation’]：从自车到全局坐标的转换（1x3 列表）。
    • info[‘sweeps’][i][‘ego2global_rotation’]：从自车到全局坐标的旋转（四元数格式的 1x4 列表）。
    • info[‘sweeps’][i][‘timestamp’]：扫描数据的时间戳。
    • info[‘sweeps’][i][‘sensor2lidar_translation’]：从当前传感器（用于收集扫描数据）到激光雷达的转换（1x3 列表）。
    • info[‘sweeps’][i][‘sensor2lidar_rotation’]：从当前传感器（用于收集扫描数据）到激光雷达的旋转（四元数格式的 1x4 列表）。
  • info[‘cams’]：相机校准信息。它包含与每个摄像头对应的六个键值： ‘CAM_FRONT’, ‘CAM_FRONT_RIGHT’, ‘CAM_FRONT_LEFT’, ‘CAM_BACK’, ‘CAM_BACK_LEFT’, ‘CAM_BACK_RIGHT’。
  • info[‘lidar2ego_translation’]：从激光雷达到自车的转换（1x3 列表）。
  • info[‘lidar2ego_rotation’]：从激光雷达到自车的旋转（四元数格式的 1x4 列表）。
  • info[‘ego2global_translation’]：从自车到全局坐标的转换（1x3 列表）。
  • info[‘ego2global_rotation’]：从自我车辆到全局坐标的旋转（四元数格式的 1x4 列表）。
  • info[‘timestamp’]：样本数据的时间戳。
  • info[‘gt_boxes’]：7 个自由度的 3D 包围框，一个 Nx7 数组。
  • info[‘gt_names’]：3D 包围框的类别，一个 1xN 数组。
  • info[‘gt_velocity’]：3D 包围框的速度（由于不准确，没有垂直测量），一个 Nx2 数组。
  • info[‘num_lidar_pts’]：每个 3D 包围框中包含的激光雷达点数。
  • info[‘num_radar_pts’]：每个 3D 包围框中包含的雷达点数。
  • info[‘valid_flag’]：每个包围框是否有效。一般情况下，我们只将包含至少一个激光雷达或毫米波雷达点的 3D 框作为有效框