BDD100K数据集图像分割实战：从数据预处理到模型训练全流程解析

1次阅读

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

BDD100K 是伯克利大学发布的自动驾驶场景数据集，包含 10 万张高清道路图像，涵盖不同天气、光照和道路条件。其语义分割标注精细划分了 40+ 类别（如道路、车辆、行人等），是训练驾驶场景理解模型的黄金标准。

新手常见痛点：

数据规模大（解压后约 1.8TB），需要高效加载策略
标注采用 JSON 格式而非常规 mask 图片，需特殊解析
类别极度不平衡（天空 / 道路占比远高于交通标志等）

数据集目录结构如下：

BDD100K/
├── images/            # 原始图片（100,000 张）│   ├── train/
│   ├── val/
├── labels/           # JSON 标注文件
│   ├── sem_seg/
│   │   ├── train/
│   │   ├── val/

使用 Python 解析 JSON 标注示例：

import json
import numpy as np
from PIL import Image

# 加载标注文件
def load_annotation(json_path):
    with open(json_path) as f:
        data = json.load(f)
    return data['frames'][0]['objects']  # 获取第一帧的所有物体标注

# 示例：查看某个标注的类别和多边形坐标
annotations = load_annotation("labels/sem_seg/train/0a0a0b1a-7c39d841.json")
for obj in annotations[:3]:  # 打印前三个物体
    print(f"类别: {obj['category']}, 多边形点数: {len(obj['poly2d'][0]['vertices'])}")

将 JSON 标注转换为二值 mask（以道路类别为例）：

def json_to_mask(json_path, img_size=(720, 1280)):
    mask = np.zeros(img_size, dtype=np.uint8)
    annotations = load_annotation(json_path)

    for obj in annotations:
        if obj['category'] == 'road':
            vertices = obj['poly2d'][0]['vertices']
            # 将多边形填充为 255（OpenCV 格式）cv2.fillPoly(mask, [np.array(vertices, np.int32)], 255)
    return mask

# 保存 mask 图像
mask = json_to_mask("sample_annotation.json")
Image.fromarray(mask).save("road_mask.png")

使用 PyTorch 的 Dataset 类处理数据：

from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T

class BDD100KDataset(Dataset):
    def __init__(self, img_dir, label_dir, transform=None):
        self.img_paths = sorted(glob(f"{img_dir}/*.jpg"))
        self.label_paths = sorted(glob(f"{label_dir}/*.json"))
        self.transform = transform

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

    def __getitem__(self, idx):
        img = Image.open(self.img_paths[idx]).convert('RGB')
        mask = json_to_mask(self.label_paths[idx])  # 调用之前的转换函数

        if self.transform:
            img = self.transform(img)
            mask = self.transform(mask)  # 相同变换保证对齐

        return img, mask

# 示例使用
transform = T.Compose([T.Resize((256, 256)),
    T.ToTensor(),])

dataset = BDD100KDataset("images/train", "labels/sem_seg/train", transform)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

实现一个轻量版 UNet：

import torch
import torch.nn as nn

class DoubleConv(nn.Module):
    """(卷积 => [BN] => ReLU) * 2"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.double_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.double_conv(x)

class UNet(nn.Module):
    def __init__(self, n_classes=40):
        super().__init__()
        # 编码器部分
        self.inc = DoubleConv(3, 64)
        self.down1 = DoubleConv(64, 128)
        # ... 完整结构参考经典 UNet

        # 解码器部分
        self.up1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        # ...

        self.outc = nn.Conv2d(64, n_classes, 1)

    def forward(self, x):
        # 实现标准的 UNet 前向传播
        x1 = self.inc(x)
        x2 = self.down1(x1)
        # ...
        return self.outc(x_final)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = UNet().to(device)
criterion = nn.CrossEntropyLoss(weight=class_weights)  # 处理类别不平衡
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(10):
    model.train()
    for images, masks in dataloader:
        images, masks = images.to(device), masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks.long())
        loss.backward()
        optimizer.step()

    # 每个 epoch 后在验证集上测试
    val_loss = evaluate(model, val_loader, device)
    print(f"Epoch {epoch}, Train Loss: {loss.item():.4f}, Val Loss: {val_loss:.4f}")

内存不足解决方案
使用 torch.utils.data.DataLoader 的num_workers参数加速数据加载
在 __getitem__ 中动态生成 mask 而非预存所有 mask
使用 torch.cuda.empty_cache() 定期清理显存

标注错误处理

# 检查标注是否包含无效多边形
for obj in annotations:
    if not obj['poly2d']:
        print(f"空标注文件: {json_path}")
        continue

训练常见问题
如果 Loss 不下降：尝试减小学习率或增加 batch size
模型预测全为背景：检查 class_weights 是否合理
GPU 利用率低：优化数据加载流程（使用 SSD 硬盘更佳）

类别	IoU	像素准确率
道路	0.912	0.956
车辆	0.763	0.881
行人	0.521	0.672

# 显示预测结果
with torch.no_grad():
    pred = model(sample_img.unsqueeze(0).to(device))
    pred_mask = torch.argmax(pred, dim=1).cpu().squeeze()

plt.figure(figsize=(12,4))
plt.subplot(1,3,1); plt.imshow(sample_img.permute(1,2,0))
plt.subplot(1,3,2); plt.imshow(true_mask)
plt.subplot(1,3,3); plt.imshow(pred_mask)