BDD100K图像分割数据集实战指南：从数据加载到模型训练

1次阅读

共计 2592 个字符，预计需要花费 7 分钟才能阅读完成。

BDD100K 是伯克利大学发布的自动驾驶场景大规模数据集，包含 10 万张高清图像，涵盖昼夜、天气、道路类型等多样化场景。其语义分割标注精细度（40+ 类别）和场景复杂性使其成为训练鲁棒模型的理想选择。但开发者常遇到两大痛点：

数据加载效率低 ：单张图像平均 5MB，传统逐图读取方式导致 I / O 瓶颈
标注格式差异 ：原生 JSON 结构需转换才能适配主流框架（如 MMSegmentation）

BDD100K/
├── images/               # 原始图像
│   ├── train/
│   ├── val/
│   └── test/
└── labels/              # 标注文件
    ├── sem_seg/
    │   ├── polygons/    # JSON 格式多边形标注
    │   └── masks/       # 官方生成的 PNG 掩码
    └── colormaps.png    # 类别颜色对照表

{
  "name": "0a0a0b1a-7c39d841.jpg",
  "attributes": {"timeofday": "daytime"},
  "labels": [
    {
      "category": "road",
      "poly2d": [{"vertices": [[x1,y1],[x2,y2],...], "types": "LLL..." }]
    }
  ]
}

import cv2
import mmap
import numpy as np

class BDDLoader:
    def __init__(self, img_dir):
        self.img_paths = sorted(Path(img_dir).glob('*.jpg'))

    def __getitem__(self, idx):
        with open(self.img_paths[idx], 'rb') as f:
            # 使用内存映射减少拷贝
            buf = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
            arr = np.frombuffer(buf, dtype=np.uint8)
            return cv2.imdecode(arr, cv2.IMREAD_COLOR)

from torch.utils.data import DataLoader
from multiprocessing import cpu_count

train_loader = DataLoader(
    dataset,
    batch_size=16,
    num_workers=cpu_count()//2,  # 留出 CPU 给模型计算
    pin_memory=True,             # 加速 GPU 传输
    persistent_workers=True
)

def convert_to_coco(json_files, output_path):
    coco = {"categories": [{"id": 1, "name": "road"}, ...],
        "images": [],
        "annotations": []}

    for img_id, ann in enumerate(json_files):
        # 处理每个多边形区域
        for obj in ann["labels"]:
            segmentation = []
            for poly in obj["poly2d"]:
                segmentation.extend([v for xy in poly["vertices"] for v in xy])

            coco["annotations"].append({"id": len(coco["annotations"]),
                "image_id": img_id,
                "category_id": class2id[obj["category"]],
                "segmentation": [segmentation],
                "area": cv2.contourArea(np.array(poly["vertices"], dtype=np.int32))
            })

    with open(output_path, 'w') as f:
        json.dump(coco, f)

import torch.nn as nn

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)

# 完整 UNet 结构定义省略...

train_transform = A.Compose([A.RandomCrop(512, 512),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])