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中。