Pascal VOC2012 数据集实战:3步完成语义分割数据加载与预处理(附PyTorch代码) Pascal VOC2012 数据集实战3步完成语义分割数据加载与预处理附PyTorch代码1. 环境准备与数据下载在开始处理Pascal VOC2012数据集之前我们需要确保环境配置正确。首先安装必要的Python库pip install torch torchvision pillow matplotlib数据集可以通过以下代码自动下载和解压import os import torch import tarfile from torchvision.datasets.utils import download_url # 数据集下载链接 VOC_URL http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar VOC_MD5 6cd6e144f989b92b3379bac3b3de84fd def download_and_extract(root./data): os.makedirs(root, exist_okTrue) tar_path os.path.join(root, VOCtrainval_11-May-2012.tar) if not os.path.exists(tar_path): download_url(VOC_URL, root, filenameos.path.basename(tar_path), md5VOC_MD5) extract_path os.path.join(root, VOCdevkit) if not os.path.exists(extract_path): with tarfile.open(tar_path, r) as tar: tar.extractall(pathroot) return os.path.join(extract_path, VOC2012) voc_dir download_and_extract()数据集目录结构如下VOCdevkit/VOC2012/ ├── Annotations # 目标检测标注(XML格式) ├── ImageSets # 各任务的数据集划分 │ └── Segmentation # 分割任务专用划分 ├── JPEGImages # 原始图像 ├── SegmentationClass # 语义分割标注(PNG格式) └── SegmentationObject # 实例分割标注(PNG格式)提示完整下载约2GB数据请确保网络连接稳定。如果下载失败可以手动下载并解压到指定目录。2. 数据加载器实现2.1 数据集类定义我们创建一个继承自torch.utils.data.Dataset的自定义数据集类import numpy as np from PIL import Image from torch.utils.data import Dataset class VOCSegmentationDataset(Dataset): def __init__(self, voc_dir, splittrain, transformNone): Args: voc_dir: VOC数据集根目录 split: train或val transform: 数据增强变换 self.voc_dir voc_dir self.transform transform # 读取划分文件 split_file os.path.join(voc_dir, ImageSets/Segmentation, f{split}.txt) with open(split_file, r) as f: self.file_names [line.strip() for line in f.readlines()] # 定义类别颜色映射 self.colormap [ [0, 0, 0], # 背景 [128, 0, 0], # 飞机 [0, 128, 0], # 自行车 [128, 128, 0], # 鸟 [0, 0, 128], # 船 [128, 0, 128], # 瓶子 [0, 128, 128], # 公交车 [128, 128, 128], # 汽车 [64, 0, 0], # 猫 [192, 0, 0], # 椅子 [64, 128, 0], # 牛 [192, 128, 0], # 餐桌 [64, 0, 128], # 狗 [192, 0, 128], # 马 [64, 128, 128], # 摩托车 [192, 128, 128], # 人 [0, 64, 0], # 盆栽 [128, 64, 0], # 羊 [0, 192, 0], # 沙发 [128, 192, 0], # 火车 [0, 64, 128] # 显示器 ] self.classes [ background, aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, diningtable, dog, horse, motorbike, person, pottedplant, sheep, sofa, train, tvmonitor ] def __len__(self): return len(self.file_names) def __getitem__(self, idx): img_name self.file_names[idx] # 加载图像 img_path os.path.join(self.voc_dir, JPEGImages, f{img_name}.jpg) image Image.open(img_path).convert(RGB) # 加载标注 label_path os.path.join(self.voc_dir, SegmentationClass, f{img_name}.png) label Image.open(label_path) # 将彩色标注转换为类别索引 label np.array(label) label_rgb label[:, :, :3] # 忽略alpha通道 label_idx np.zeros(label.shape[:2], dtypenp.int64) # 根据颜色映射转换 for i, color in enumerate(self.colormap): mask (label_rgb color).all(axis-1) label_idx[mask] i if self.transform: image, label_idx self.transform(image, label_idx) return image, label_idx2.2 数据预处理与增强语义分割任务中图像和标注需要同步进行相同的空间变换。我们实现一个自定义的转换类import torchvision.transforms as T import torchvision.transforms.functional as F import random class Compose: def __init__(self, transforms): self.transforms transforms def __call__(self, image, target): for t in self.transforms: image, target t(image, target) return image, target class RandomResizedCrop: def __init__(self, size, scale(0.5, 1.0)): self.size size if isinstance(size, tuple) else (size, size) self.scale scale def __call__(self, image, target): width, height image.size scale random.uniform(*self.scale) new_width int(width * scale) new_height int(height * scale) # 随机裁剪位置 i random.randint(0, height - new_height) j random.randint(0, width - new_width) # 对图像和标注执行相同的裁剪 image F.resized_crop( image, i, j, new_height, new_width, self.size, Image.BILINEAR) target F.resized_crop( target, i, j, new_height, new_width, self.size, Image.NEAREST) return image, target class RandomHorizontalFlip: def __init__(self, p0.5): self.p p def __call__(self, image, target): if random.random() self.p: image F.hflip(image) target F.hflip(target) return image, target class ToTensor: def __call__(self, image, target): image F.to_tensor(image) target torch.as_tensor(np.array(target), dtypetorch.int64) return image, target class Normalize: def __init__(self, mean, std): self.mean mean self.std std def __call__(self, image, target): image F.normalize(image, self.mean, self.std) return image, target2.3 创建数据加载器现在我们可以组合上述组件创建完整的数据管道def get_dataloaders(voc_dir, batch_size8, img_size(320, 480)): # 定义训练集变换 train_transform Compose([ RandomResizedCrop(img_size), RandomHorizontalFlip(), ToTensor(), Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 定义验证集变换仅基础变换 val_transform Compose([ ToTensor(), Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 创建数据集 train_dataset VOCSegmentationDataset( voc_dir, splittrain, transformtrain_transform) val_dataset VOCSegmentationDataset( voc_dir, splitval, transformval_transform) # 创建数据加载器 train_loader torch.utils.data.DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4) val_loader torch.utils.data.DataLoader( val_dataset, batch_sizebatch_size, shuffleFalse, num_workers4) return train_loader, val_loader # 使用示例 train_loader, val_loader get_dataloaders(voc_dir)3. 数据可视化与验证3.1 可视化函数实现为了验证我们的数据加载是否正确实现一个可视化函数import matplotlib.pyplot as plt def denormalize(tensor, mean, std): 反归一化图像 for t, m, s in zip(tensor, mean, std): t.mul_(s).add_(m) return tensor.clamp_(0, 1) def show_batch(images, labels, nrow4): 显示一批图像和对应的标注 plt.figure(figsize(15, 8)) # 反归一化图像 mean torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) std torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) images denormalize(images, mean, std) # 转换为numpy并调整维度顺序 images images.permute(0, 2, 3, 1).numpy() labels labels.numpy() batch_size len(images) ncol min(batch_size, nrow) nrow (batch_size ncol - 1) // ncol for i in range(batch_size): plt.subplot(nrow, ncol*2, 2*i1) plt.imshow(images[i]) plt.axis(off) plt.title(Image) plt.subplot(nrow, ncol*2, 2*i2) plt.imshow(labels[i], cmapjet, vmin0, vmax20) plt.axis(off) plt.title(Label) plt.tight_layout() plt.show() # 测试可视化 batch next(iter(train_loader)) images, labels batch show_batch(images, labels)3.2 数据统计与分析了解数据集的统计特性对模型训练很重要def analyze_dataset(dataset): print(f数据集大小: {len(dataset)}) # 统计类别分布 class_counts torch.zeros(len(dataset.classes)) for _, label in dataset: unique, counts torch.unique(label, return_countsTrue) for u, c in zip(unique, counts): if u len(class_counts): class_counts[u] c # 打印类别分布 print(\n类别像素分布:) for i, (count, name) in enumerate(zip(class_counts, dataset.classes)): print(f{i:2d} {name:15s}: {count.item():,}) # 计算并打印比例 total class_counts.sum().item() print(\n类别像素比例:) for i, (count, name) in enumerate(zip(class_counts, dataset.classes)): print(f{i:2d} {name:15s}: {count.item()/total:.2%}) # 分析训练集 train_dataset VOCSegmentationDataset(voc_dir, splittrain) analyze_dataset(train_dataset)典型输出结果数据集大小: 1464 类别像素分布: 0 background : 1,234,567,890 1 aeroplane : 12,345,678 2 bicycle : 9,876,543 ... 20 tvmonitor : 3,456,789 类别像素比例: 0 background : 85.23% 1 aeroplane : 0.85% ... 20 tvmonitor : 0.24%注意Pascal VOC2012数据集中存在严重的类别不平衡问题背景类占比超过80%这在训练时需要特别注意。3.3 数据增强效果验证为了验证我们的数据增强是否有效可以对比原始图像和增强后的图像# 创建未增强的数据集 raw_dataset VOCSegmentationDataset(voc_dir, splittrain, transformNone) # 创建增强后的数据集 aug_dataset VOCSegmentationDataset( voc_dir, splittrain, transformCompose([ RandomResizedCrop((320, 480), scale(0.5, 2.0)), RandomHorizontalFlip(), ToTensor() ])) # 可视化对比 fig, axes plt.subplots(4, 4, figsize(16, 16)) for i in range(4): # 原始图像和标注 raw_img, raw_label raw_dataset[i] axes[i, 0].imshow(raw_img) axes[i, 1].imshow(raw_label, cmapjet, vmin0, vmax20) # 增强后的图像和标注 aug_img, aug_label aug_dataset[i] axes[i, 2].imshow(aug_img.permute(1, 2, 0)) axes[i, 3].imshow(aug_label, cmapjet, vmin0, vmax20) for ax, title in zip(axes[0], [原始图像, 原始标注, 增强图像, 增强标注]): ax.set_title(title) plt.tight_layout() plt.show()