### UNet 图像分割 PyTorch 实现 CamVid 数据集 代码示例 对于CamVid数据集上的UNet图像分割任务，可以采用如下PyTorch实现方案。此实现基于`segmentation-models-pytorch`库来构建模型并训练。 #### 安装必要的依赖包 为了顺利运行下面的代码片段，需要先安装所需的Python库[^1]: ```bash pip install segmentation-models-pytorch torch torchvision ``` #### 导入所需模块 加载必需的Python库和工具： ```python import os from PIL import Image import numpy as np import albumentations as A from albumentations.pytorch.transforms import ToTensorV2 import torch from torch.utils.data import Dataset, DataLoader import segmentation_models_pytorch as smp ``` #### 加载与预处理数据 定义自定义的数据集类以适应CamVid数据结构，并应用适当的数据增强技术提高泛化能力： ```python class CamVidDataset(Dataset): def __init__(self, image_dir, mask_dir, transform=None): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.images = os.listdir(image_dir) def __len__(self): return len(self.images) def __getitem__(self, index): img_path = os.path.join(self.image_dir, self.images[index]) mask_path = os.path.join(self.mask_dir, self.images[index].replace('.png', '_L.png')) image = np.array(Image.open(img_path).convert("RGB")) mask = np.array(Image.open(mask_path).convert("L"), dtype=np.float32) mask[mask == 255.0] = 1.0 if self.transform is not None: augmentations = self.transform(image=image, mask=mask) image = augmentations["image"] mask = augmentations["mask"] return image, mask ``` #### 创建转换函数 设置训练期间使用的变换操作，包括随机裁剪、水平翻转等： ```python transform = A.Compose( [ A.Resize(height=224, width=224), A.HorizontalFlip(p=0.5), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ] ) ``` #### 初始化数据加载器 准备用于迭代访问整个数据集的对象： ```python train_dataset = CamVidDataset( image_dir="data/CamVid/train", mask_dir="data/CamVid/train_labels", transform=transform, ) val_dataset = CamVidDataset( image_dir="data/CamVid/val", mask_dir="data/CamVid/val_labels", transform=transform, ) train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=2) val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=2) ``` #### 构建U-Net模型 利用`segmentation-models-pytorch`创建一个带有ResNet34编码器的U-Net架构实例： ```python model = smp.Unet(encoder_name='resnet34', encoder_weights='imagenet') device = 'cuda' if torch.cuda.is_available() else 'cpu' model.to(device); ``` #### 训练过程概览 虽然完整的训练循环超出了当前讨论范围，但通常涉及定义损失函数（如交叉熵），优化算法（如AdamW），以及通过多个epoch遍历数据集更新权重的过程。