📚 课程目录第2章 BEV与Occupancy数据集准备与预处理

2. BEV与Occupancy数据集准备与预处理

任务目标

本案例完成BEV与Occupancy联合训练前的数据准备闭环：下载nuScenes mini数据集、解压并转换为标准格式，随后实现相机内外参标定与图像去畸变预处理，输出可直接用于后续模型训练的张量数据。

环境依赖

需要安装：nuscenes-devkit、opencv-python、numpy、pyquaternion。数据集下载地址：https://www.nuscenes.org/download（选择mini版本，约4GB）。

数据下载与格式转换

下载完成后，使用以下脚本完成解压与nuScenes标准格式的初始化。该步骤将原始文件组织为数据库表结构，并建立相机与激光雷达的时间戳对齐索引。

import os
import tarfile
from nuscenes import NuScenes

# 1. 解压数据集
def extract_nuscenes(tar_path, extract_dir):
    os.makedirs(extract_dir, exist_ok=True)
    with tarfile.open(tar_path, 'r') as tar:
        tar.extractall(extract_dir)
    print(f"解压完成: {extract_dir}")

# 2. 初始化nuScenes数据库（生成samples、sample_data等表）
def init_nuscenes(dataroot, version='v1.0-mini'):
    nusc = NuScenes(version=version, dataroot=dataroot, verbose=True)
    return nusc

if __name__ == '__main__':
    # 请根据实际下载路径修改
    tar_path = './nuScenes-mini.tgz'
    dataroot = './nuscenes-mini'
    extract_nuscenes(tar_path, dataroot)
    nusc = init_nuscenes(dataroot)
    print(f"共加载 {len(nusc.sample)} 个关键帧样本")

相机标定参数提取与去畸变预处理

核心预处理流程：从nuScenes数据库读取每个相机的外参（ego位姿）与内参（相机标定），利用OpenCV对原始图像进行去畸变校正，并生成校正后的图像与对应的投影矩阵。

import numpy as np
import cv2
from pyquaternion import Quaternion
from nuscenes.utils.data_classes import Box
from nuscenes.utils.geometry_utils import transform_matrix

def load_calibration(nusc, sample_data_record):
    """提取相机内外参"""
    # 内参
    cam_data = nusc.get('sample_data', sample_data_record[''])
    cam_sensor = nusc.get('calibrated_sensor', cam_data['calibrated_sensor_'])
    K = np.array(cam_sensor['camera_intrinsic'])  # 3x3内参矩阵
    
    # 外参：传感器到ego
    sensor_to_ego = transform_matrix(
        cam_sensor['translation'], 
        Quaternion(cam_sensor['rotation']), 
        inverse=False
    )
    
    # ego到全局
    ego_pose = nusc.get('ego_pose', cam_data['ego_pose_'])
    ego_to_global = transform_matrix(
        ego_pose['translation'], 
        Quaternion(ego_pose['rotation']), 
        inverse=False
    )
    
    # 合成相机到全局的外参
    cam_to_global = ego_to_global @ sensor_to_ego
    return K, cam_to_global

def undistort_image(nusc, sample_, cam_channel='CAM_FRONT'):
    """对指定相机的图像进行去畸变处理"""
    # 获取样本数据记录
    sample = nusc.get('sample', sample_)
    cam_data = nusc.get('sample_data', sample['data'][cam_channel])
    
    # 读取原始图像
    img_path = os.path.join(nusc.dataroot, cam_data['filename'])
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    
    # 获取标定参数
    K, cam_to_global = load_calibration(nusc, cam_data)
    
    # 去畸变（nuScenes使用pinhole模型，畸变系数为0）
    # 实际工程中若存在畸变，需从calibrated_sensor中读取distortion系数
    dist_coeffs = np.zeros(5)  # nuScenes mini默认无畸变
    new_K, roi = cv2.getOptimalNewCameraMatrix(K, dist_coeffs, (w, h), 1, (w, h))
    
    # 校正映射
    mapx, mapy = cv2.initUndistortRectifyMap(
        K, dist_coeffs, None, new_K, (w, h), cv2.CV_32FC1
    )
    undistorted_img = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)
    
    # 计算校正后的投影矩阵（用于BEV投影）
    R = np.eye(3)  # 无旋转校正
    P = new_K @ np.hstack((R, np.zeros((3, 1))))  # 3x4投影矩阵
    
    return undistorted_img, new_K, cam_to_global, P

# 批量处理示例：处理第一个样本的所有6个相机
def preprocess_sample(nusc, sample_):
    cam_channels = ['CAM_FRONT', 'CAM_FRONT_LEFT', 'CAM_FRONT_RIGHT',
                    'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_BACK_RIGHT']
    processed = {}
    for cam in cam_channels:
        img, K, T, P = undistort_image(nusc, sample_, cam)
        processed[cam] = {
            'image': img,
            'K': K,          # 校正后内参
            'T': T,          # 相机到全局的变换矩阵
            'P': P           # 投影矩阵
        }
    return processed

if __name__ == '__main__':
    # 使用之前初始化的nusc对象
    sample_ = nusc.sample[0]['']
    result = preprocess_sample(nusc, sample_)
    
    # 验证输出形状
    for cam, data in result.items():
        print(f"{cam}: 图像尺寸 {data['image'].shape}, "
              f"内参形状 {data['K'].shape}, "
              f"外参形状 {data['T'].shape}")

运行与验证

执行步骤：

将第一个脚本保存为 download_and_init.py，第二个脚本保存为 preprocess.py，依次运行。
运行前需在preprocess.py中导入第一个脚本生成的nusc对象（或合并执行）。

验证通过条件：

控制台输出共加载 404 个关键帧样本（mini版样本数），且无异常报错。
预处理脚本输出6个相机的图像尺寸均为 (900, 1600, 3)，内参矩阵形状为 (3, 3)，外参矩阵为 (4, 4)。
随机选取一张校正后图像，与原始图像对比，视觉上无畸变扭曲（本数据集无畸变，校正前后应基本一致）。

常见问题

下载速度慢：建议使用官方提供的nuScenes下载脚本，支持断点续传。
内存不足：mini数据集约4GB，需保证至少8GB可用内存；若处理全量数据集，建议分批加载。
时间戳对齐：本案例使用sample的data字段直接获取同步好的相机数据，无需额外处理时间同步。