Cityscapes数据集实战指南：从数据加载到模型训练的全流程解析

1次阅读

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

Cityscapes 数据集包含来自 50 个不同城市的街道场景图像，主要用于语义分割任务。数据集的核心特点包括高分辨率（2048×1024）、精细标注（fine annotation）和粗糙标注（coarse annotation）两种标注类型，以及 19 类语义标签。

数据组织方式：
leftImg8bit：原始 RGB 图像
gtFine：精细标注（2975 张）
gtCoarse：粗糙标注（19998 张）
19 类语义标签（按 trainId 顺序）：
0: 可忽略区域
1: 道路
2: 人行道
3: 建筑
4: 围墙
5: 栅栏
6: 电线杆
7: 交通灯
8: 交通标志
9: 植被
10: 地形
11: 天空
12: 行人
13: 骑行者
14: 小车
15: 卡车
16: 巴士
17: 火车
18: 摩托车

标注文件解析难点：
精细标注使用 labelIds 格式（每个像素对应一个类别 ID）
粗糙标注的标注质量较低，建议优先使用精细标注
解决方案：使用官方提供的 cityscapesscripts 工具包
内存管理问题：
高分辨率图像直接加载会消耗大量内存
解决方案：
- 使用 PIL.Image 的懒加载模式
- 实现按需裁剪的数据流
标注对齐问题：
原始图像和标注文件需要严格匹配
解决方案：维护文件名映射表

import os
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms

class CityscapesDataset(Dataset):
    """Cityscapes 数据集加载器"""

    def __init__(self, root_dir, split='train', crop_size=(512, 1024)):
        self.root_dir = root_dir
        self.split = split
        self.crop_size = crop_size

        # 文件路径收集
        self.image_paths = []
        self.label_paths = []

        cities = os.listdir(os.path.join(root_dir, 'leftImg8bit', split))
        for city in cities:
            img_dir = os.path.join(root_dir, 'leftImg8bit', split, city)
            label_dir = os.path.join(root_dir, 'gtFine', split, city)

            for img_name in os.listdir(img_dir):
                if img_name.endswith('_leftImg8bit.png'):
                    base_name = img_name.replace('_leftImg8bit.png', '')
                    label_name = base_name + '_gtFine_labelIds.png'

                    self.image_paths.append(os.path.join(img_dir, img_name))
                    self.label_paths.append(os.path.join(label_dir, label_name))

        # 数据增强
        self.transform = transforms.Compose([transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5)
        ])

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

    def __getitem__(self, idx):
        # 加载图像和标注
        image = Image.open(self.image_paths[idx]).convert('RGB')
        label = Image.open(self.label_paths[idx])

        # 随机裁剪
        i, j, h, w = transforms.RandomCrop.get_params(image, output_size=self.crop_size)
        image = transforms.functional.crop(image, i, j, h, w)
        label = transforms.functional.crop(label, i, j, h, w)

        # 数据增强（仅对图像）image = self.transform(image)

        # 转换为 Tensor
        image = transforms.ToTensor()(image)
        label = torch.from_numpy(np.array(label)).long()

        return image, label

基础训练流程：

import torch.nn as nn
from torch.utils.data import DataLoader

# 初始化数据集
dataset = CityscapesDataset('path/to/cityscapes', split='train')
dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=4)

# 定义模型（示例使用简化版 UNet）model = UNet(num_classes=19).cuda()
criterion = nn.CrossEntropyLoss(ignore_index=255)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 训练循环
for epoch in range(100):
    for images, labels in dataloader:
        images = images.cuda()
        labels = labels.cuda()

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()