告别像素级标注！用PyTorch和CAM实现图像级标签的语义分割（附完整代码）

告别像素级标注！用PyTorch和CAM实现图像级标签的语义分割实战指南

在计算机视觉领域，语义分割一直扮演着至关重要的角色。它能精确到像素级别识别图像中的每个对象，为自动驾驶、医疗影像分析等场景提供关键技术支持。然而，获取像素级标注数据的过程往往耗时耗力——标注一张Cityscapes数据集中的图像平均需要90分钟，而医学图像标注成本更高达数百美元每张。这种高昂的标注成本成为许多项目难以跨越的门槛。

弱监督语义分割(WSSS)技术应运而生，它让我们能够利用更易获得的图像级标签（即整图分类标签）来训练分割模型。本文将手把手带你实现一个基于PyTorch和类激活图(CAM)的完整解决方案，仅用图像级标签就能生成高质量的分割结果。以下是我们的技术路线：

1. 使用ResNet等分类网络训练图像分类模型
2. 通过CAM提取初始分割区域
3. 应用CRF等后处理技术优化分割边界
4. 用生成的伪标签训练最终分割模型

1. 环境准备与数据加载

1.1 基础环境配置

推荐使用Python 3.8+和PyTorch 1.10+环境。以下是必需的依赖包：

pip install torch torchvision opencv-python pillow pandas scikit-image pycocotools

对于GPU加速，建议安装对应版本的CUDA工具包。可以通过以下命令验证PyTorch是否正确识别了GPU：

import torch print(torch.cuda.is_available()) # 应输出True print(torch.__version__) # 确认版本号

1.2 数据集处理

我们以PASCAL VOC 2012数据集为例，它同时包含图像级标签和像素级标注（仅用于验证）。数据目录结构应如下：

VOC2012/ ├── JPEGImages/ # 原始图像 ├── ImageSets/ │ └── Segmentation/ # 划分文件 └── SegmentationClass/ # 真实标注（仅测试用）

实现自定义数据集类时，关键要处理好图像变换和标签编码：

from torch.utils.data import Dataset import torchvision.transforms as T class VOCDataset(Dataset): def __init__(self, img_dir, label_file, transform=None): self.img_dir = img_dir self.labels = pd.read_csv(label_file) # 图像名,类别1,类别2... self.transform = transform or T.Compose([ T.Resize(256), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.labels.iloc[idx, 0]+'.jpg') image = Image.open(img_path).convert('RGB') label = torch.tensor(self.labels.iloc[idx, 1:].values.astype('float32')) return self.transform(image), label

提示：对于多标签分类任务，使用sigmoid激活而非softmax，损失函数应选择BCEWithLogitsLoss

2. 分类模型训练与CAM生成

2.1 修改ResNet获取特征图

标准的分类网络会丢弃空间信息，我们需要修改最后几层结构以保留特征图：

import torch.nn as nn from torchvision.models import resnet50 class CAMResNet(nn.Module): def __init__(self, num_classes): super().__init__() base = resnet50(pretrained=True) self.features = nn.Sequential(*list(base.children())[:-2]) self.gap = nn.AdaptiveAvgPool2d(1) self.fc = nn.Conv2d(2048, num_classes, kernel_size=1) def forward(self, x): features = self.features(x) # [B,2048,H,W] logits = self.fc(features) # [B,C,H,W] pooled = self.gap(logits) # [B,C,1,1] return pooled.squeeze(), logits

关键修改点：

• 移除原ResNet的全局平均池化层和全连接层
• 添加1x1卷积层直接输出类别数通道的特征图
• 同时返回全局预测结果和空间特征图

2.2 训练分类模型

训练过程与常规分类任务类似，但要注意多标签场景的特殊处理：

model = CAMResNet(num_classes=20).cuda() criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9) for epoch in range(50): for inputs, labels in train_loader: inputs, labels = inputs.cuda(), labels.cuda() optimizer.zero_grad() outputs, _ = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step()

注意：学习率策略建议使用余弦退火，batch size不宜过小（推荐≥32）

2.3 生成类激活图

训练完成后，我们可以提取每张图的CAM作为初始分割：

def generate_cam(model, img_tensor, target_class): _, logits = model(img_tensor.unsqueeze(0).cuda()) weights = model.fc.weight[target_class] # [2048] # 计算类激活图 cam = torch.matmul(weights, logits.squeeze().permute(1,2,0)) cam = cam.cpu().numpy() cam = cv2.resize(cam, (img_tensor.shape[1], img_tensor.shape[0])) cam = np.maximum(cam, 0) # ReLU cam = (cam - cam.min()) / (cam.max() - cam.min()) return cam

典型CAM结果会高亮与目标类别最相关的区域，但往往存在以下问题：

• 仅激活最具判别性的部分（如狗头而非全身）
• 边界粗糙，缺乏精确轮廓
• 对小物体响应较弱

3. 伪标签优化策略

3.1 基于CRF的后处理

条件随机场(CRF)能有效改善CAM的粗糙边界。以下是使用OpenCV实现的DenseCRF：

import pydensecrf.densecrf as dcrf from pydensecrf.utils import unary_from_softmax def apply_crf(img, cam, n_classes=2): h, w = img.shape[:2] # 准备一元势能 probs = np.stack([1-cam, cam], axis=0) U = unary_from_softmax(probs) # 创建CRF d = dcrf.DenseCRF2D(w, h, n_classes) d.setUnaryEnergy(U) # 添加二元势能 d.addPairwiseGaussian(sxy=3, compat=3) d.addPairwiseBilateral(sxy=20, srgb=3, rgbim=img, compat=10) # 推理 Q = d.inference(5) return np.argmax(Q, axis=0).reshape(h, w)

关键参数说明：

• sxy：空间高斯核的标准差
• srgb：颜色高斯核的标准差
• compat：兼容性系数（越大越平滑）

3.2 基于AffinityNet的改进

CRF主要处理局部关系，而AffinityNet能学习像素间的长距离依赖。实现步骤如下：

1. 使用CAM生成初始种子区域
2. 训练AffinityNet预测像素间相似度
3. 通过随机游走传播标签

class AffinityNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.fc = nn.Conv2d(128, 1, kernel_size=1) def forward(self, x1, x2): h1 = F.relu(self.conv1(x1)) h2 = F.relu(self.conv1(x2)) h = torch.cat([h1, h2], dim=1) return torch.sigmoid(self.fc(h))

训练AffinityNet时，正样本来自CAM高响应区域内的像素对，负样本来自高响应与低响应区域间的像素对。

4. 完整训练流程与结果评估

4.1 端到端训练方案

结合前述组件，完整的弱监督训练流程如下表所示：

4.2 分割模型实现

我们选择DeepLabv3+作为最终分割模型，其多尺度特征融合能力适合处理CAM生成的不完整标注：

from torchvision.models.segmentation import deeplabv3_resnet50 model = deeplabv3_resnet50(pretrained=False, num_classes=21) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) for epoch in range(100): for img, pseudo_label in dataloader: outputs = model(img)['out'] loss = dice_loss(outputs, pseudo_label) optimizer.zero_grad() loss.backward() optimizer.step()

提示：使用Dice损失而非交叉熵，对不均衡标注更鲁棒

4.3 性能评估指标

在PASCAL VOC验证集上的典型结果：

可视化对比显示，我们的方法虽然略逊于全监督模型，但显著优于原始CAM结果，特别是在物体边界和小目标处理上。