深度学习实战：基于Oxford-IIIT Pet Dataset的宠物图像分割与分类

1. 从零开始玩转宠物图像分割第一次接触Oxford-IIIT Pet Dataset时我被这个数据集的精细程度惊艳到了。作为计算机视觉领域的经典数据集它包含了7349张猫狗图片每张都配有像素级的分割掩码。简单来说就是能精确到像素级别区分宠物身体的各个部位比如耳朵、爪子、尾巴等。这对于想入门图像分割的新手来说简直是宝藏资源。我建议初学者先从数据探索开始。解压后的数据集通常包含images和annotations两个核心文件夹。images里是各种可爱的宠物照片而annotations/trimaps下的黑白掩码图才是重点。这些掩码用不同灰度值标记了宠物身体部位1-前景2-背景3-轮廓。记得我第一次加载图片时犯了个低级错误——没把掩码图转为灰度模式导致模型训练完全跑偏。from PIL import Image import matplotlib.pyplot as plt # 正确加载方式示例 img Image.open(images/dog.1.jpg) mask Image.open(annotations/trimaps/dog.1.png).convert(L) # 关键步骤 plt.imshow(mask, cmapgray)2. 数据预处理中的那些坑原始数据直接扔给模型的效果往往很差这里分享几个实战中总结的预处理技巧。首先是尺寸统一化数据集中的图片尺寸不一我用CenterCropResize组合拳解决from torchvision import transforms transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor() ])更关键的是掩码处理。由于原始trimap是单通道8位图0-255需要先归一化到0-1范围再转换为类别标签。这里有个隐藏坑不同框架对类别索引的要求不同。PyTorch要求从0开始连续编号而TensorFlow可能要求one-hot编码。我常用的转换方法是import numpy as np def process_mask(mask): mask np.array(mask) mask[mask 2] 1 # 背景转1 mask[mask 3] 2 # 轮廓转2 return mask数据增强方面我发现对宠物图片最有效的是随机水平翻转色彩抖动。但要注意同步处理图像和掩码曾经因为忘记对掩码做相同变换导致模型学到错误对应关系。推荐使用Albumentations库import albumentations as A aug A.Compose([ A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.2), ], additional_targets{mask: mask})3. 模型选型与调参实战经过多次实验我发现U-Net在这个任务上表现非常突出。不同于直接套用原版我对网络结构做了三处针对性改进编码器部分改用ResNet34预训练权重在跳跃连接处添加了空间注意力模块输出层改用带边缘权重的Dice Loss具体实现时有个容易忽略的细节预训练模型的输入归一化参数需要与训练时一致。我栽过跟头后才记住ImageNet的mean[0.485,0.456,0.406]std[0.229,0.224,0.225]。import torch.nn as nn from segmentation_models_pytorch import Unet model Unet( encoder_nameresnet34, encoder_weightsimagenet, classes3, activationNone )训练策略上推荐采用分阶段学习率前5轮用1e-4预热之后切换为余弦退火。batch_size设为32时在RTX 3090上每个epoch约2分钟。验证集IoU能达到0.89的关键是加入了边缘惩罚项class EdgeAwareLoss(nn.Module): def __init__(self): super().__init__() self.sobel SobelOperator() def forward(self, pred, target): edge_mask self.sobel(target) loss (1 edge_mask) * F.cross_entropy(pred, target) return loss.mean()4. 评估与结果可视化模型评估不能只看准确率对于分割任务我主要监控三个指标mIoU平均交并比Boundary F1 Score各类别召回率可视化环节特别重要。我习惯用wandb记录训练过程这个代码片段可以生成带预测对比的网格图def plot_predictions(batch, n3): fig, axes plt.subplots(n, 3, figsize(12,4*n)) for i in range(n): axes[i,0].imshow(batch[image][i].permute(1,2,0)) axes[i,1].imshow(batch[mask][i].squeeze(), cmapjet) axes[i,2].imshow(torch.argmax(outputs[i],0).cpu(), cmapjet) return fig实际部署时遇到过一个典型问题模型在测试集表现很好但用户上传的图片效果差。原因是训练数据都是专业拍摄的正面照而真实场景图片存在遮挡、模糊等情况。解决方法是通过GAN生成更多困难样本加入训练。5. 进阶技巧与优化方向当基础模型达到瓶颈时可以尝试这些进阶方法使用swin transformer作为编码器加入OCR模块提升小物体分割效果采用CascadePSP进行后处理内存优化方面我发现混合精度训练梯度累积能大幅降低显存消耗。对于移动端部署用TensorRT量化后的模型体积能缩小到原来的1/4model.half() # 转为半精度 for param in model.parameters(): param.requires_grad False最近在尝试的主动学习策略也很有前景先用少量数据训练初始模型然后让模型自动筛选信息量大的未标注样本进行人工标注。实验表明这种方法能减少30%的标注成本。