4. BEV与Occupancy数据加载器构建

任务目标

构建一个用于BEV与Occupancy联合训练的数据加载器，支持多线程读取、批处理、数据打乱，并确保BEV标签与Occupancy标签在采样时保持同步。本案例以NuScenes格式的mini split为示例数据源，实现一个最小可复用的BEVOccDataset类。

环境与依赖

Python ≥ 3.8，PyTorch ≥ 1.10，numpy，pyquaternion（用于旋转矩阵）
数据目录结构（本案例自带生成脚本，无需外部下载）：


data/
├── samples/
│   ├── sample_0001.npz      # 包含 lidar_pc (N,4)、pose (4,4)
│   └── sample_0002.npz
└── annotations/
    ├── bev_label_0001.npz   # 包含 bev_map (128,128)
    ├── occ_label_0001.npz   # 包含 occ_grid (32,32,32)
    └── ...（对应同名sample）

关键完整代码


import os
import glob
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.dataloader import default_collate

class BEVOccDataset(Dataset):
    """同步加载BEV与Occupancy标签的数据集"""
    def __init__(self, data_root, split='train', shuffle_seed=0):
        self.sample_dir = os.path.join(data_root, 'samples')
        self.annot_dir = os.path.join(data_root, 'annotations')
        # 以sample文件名为唯一ID，保证BEV与Occ标签一一对应
        self.sample_ids = [os.path.basename(f).replace('.npz', '')
                           for f in glob.glob(os.path.join(self.sample_dir, '*.npz'))]
        self.sample_ids.sort()
        if shuffle_seed is not None:
            rng = np.random.RandomState(shuffle_seed)
            rng.shuffle(self.sample_ids)
        # 按比例切分train/val（80/20）
        n = len(self.sample_ids)
        split_idx = int(n * 0.8)
        if split == 'train':
            self.sample_ids = self.sample_ids[:split_idx]
        else:
            self.sample_ids = self.sample_ids[split_idx:]

    def __len__(self):
        return len(self.sample_ids)

    def __getitem__(self, idx):
        sid = self.sample_ids[idx]
        # 同步读取：同一sid下同时加载BEV和Occ标签
        sample = np.load(os.path.join(self.sample_dir, f'{sid}.npz'))
        bev = np.load(os.path.join(self.annot_dir, f'bev_label_{sid}.npz'))
        occ = np.load(os.path.join(self.annot_dir, f'occ_label_{sid}.npz'))

        lidar = sample['lidar_pc'].astype(np.float32)          # (N,4)
        pose = sample['pose'].astype(np.float32)               # (4,4)
        bev_map = bev['bev_map'].astype(np.float32)            # (128,128)
        occ_grid = occ['occ_grid'].astype(np.float32)          # (32,32,32)

        # 转换为张量（CHW格式便于后续网络处理）
        bev_tensor = torch.from_numpy(bev_map).unsqueeze(0)    # (1,128,128)
        occ_tensor = torch.from_numpy(occ_grid).unsqueeze(0)   # (1,32,32,32)
        return {
            'lidar': torch.from_numpy(lidar),
            'pose': torch.from_numpy(pose),
            'bev_label': bev_tensor,
            'occ_label': occ_tensor,
            'sample_id': sid
        }

def bev_occ_collate(batch):
    """自定义批处理：处理变长点云，BEV/Occ标签保持同步堆叠"""
    lidar_list = [b['lidar'] for b in batch]
    # 对点云进行padding到batch内最大长度
    max_len = max(l.shape[0] for l in lidar_list)
    lidar_padded = torch.zeros(len(batch), max_len, 4)
    for i, l in enumerate(lidar_list):
        lidar_padded[i, :l.shape[0]] = l
    return {
        'lidar': lidar_padded,
        'pose': torch.stack([b['pose'] for b in batch]),
        'bev_label': torch.stack([b['bev_label'] for b in batch]),
        'occ_label': torch.stack([b['occ_label'] for b in batch]),
        'sample_id': [b['sample_id'] for b in batch]
    }

def build_dataloader(data_root, batch_size=4, num_workers=2, split='train'):
    """构建多线程数据加载器"""
    dataset = BEVOccDataset(data_root, split=split)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(split == 'train'),   # 训练集打乱，验证集不打乱
        num_workers=num_workers,
        collate_fn=bev_occ_collate,
        pin_memory=True,
        drop_last=(split == 'train')
    )
    return loader

# 最小数据生成脚本（用于验证）
def generate_demo_data(data_root='./data', num_samples=10):
    os.makedirs(f'{data_root}/samples', exist_ok=True)
    os.makedirs(f'{data_root}/annotations', exist_ok=True)
    for i in range(num_samples):
        sid = f'sample_{i:04d}'
        # 随机点云（100~200点）
        lidar = np.random.randn(np.random.randint(100, 200), 4).astype(np.float32)
        pose = np.eye(4, dtype=np.float32)
        np.savez(f'{data_root}/samples/{sid}.npz', lidar_pc=lidar, pose=pose)
        # BEV标签（128x128二值图）
        bev = (np.random.rand(128, 128) > 0.7).astype(np.float32)
        np.savez(f'{data_root}/annotations/bev_label_{sid}.npz', bev_map=bev)
        # Occ标签（32x32x32二值体素）
        occ = (np.random.rand(32, 32, 32) > 0.9).astype(np.float32)
        np.savez(f'{data_root}/annotations/occ_label_{sid}.npz', occ_grid=occ)
    print(f'生成 {num_samples} 个演示样本于 {data_root}')

if __name__ == '__main__':
    # 运行验证
    generate_demo_data()
    loader = build_dataloader('./data', batch_size=4, num_workers=2, split='train')
    batch = next(iter(loader))
    print('Batch keys:', list(batch.keys()))
    print('lidar shape:', batch['lidar'].shape)       # (4, max_len, 4)
    print('bev shape:', batch['bev_label'].shape)     # (4, 1, 128, 128)
    print('occ shape:', batch['occ_label'].shape)     # (4, 1, 32, 32, 32)
    print('样本数（train）:', len(loader.dataset))

运行与验证

将上述代码保存为 bev_occ_dataloader.py，直接运行：python bev_occ_dataloader.py
预期输出（确定性结果）：


生成 10 个演示样本于 ./data
Batch keys: ['lidar', 'pose', 'bev_label', 'occ_label', 'sample_id']
lidar shape: (4, max_len, 4)
bev shape: (4, 1, 128, 128)
occ shape: (4, 1, 32, 32, 32)
样本数（train）: 8

其中max_len为当前batch内点云最大点数（介于100~200之间）。

验证判定标准

同步性验证：打印的sample_id列表中，每个ID在samples/与annotations/下均存在对应文件，且BEV与Occ标签来自同一ID。
批处理验证：bev_label与occ_label的第一维均为batch_size（4），且一一对应。
多线程验证：设置num_workers=2后程序正常迭代无报错，说明多进程读取生效。
打乱验证：连续两次运行训练集loader，sample_id顺序不同；验证集loader顺序固定。

关键说明

本案例使用npz格式存储单帧数据，实际工程中可替换为数据库或内存映射以提升IO效率。
自定义collate_fn是处理变长点云的必要组件，BEV与Occ标签为定长张量，直接stack即可保持同步。
若需在训练中同时使用多帧时序信息，可在__getitem__中按时间戳索引相邻帧，但本案例聚焦单帧同步加载，不扩展时序逻辑。

📚 课程目录

4. BEV与Occupancy数据加载器构建

任务目标

环境与依赖

关键完整代码

运行与验证

验证判定标准

关键说明