从人脸编辑到图像分割：如何用CelebAMask-HQ训练你的第一个语义分割模型（PyTorch版）

从人脸编辑到图像分割CelebAMask-HQ语义分割实战指南当我们需要训练一个能够精确识别人脸各个部位的AI模型时高质量的数据集是关键。CelebAMask-HQ作为目前最全面的人脸语义分割数据集之一包含了30,000张512×512分辨率的人脸图像每张图像都有19个面部组件的精细标注。这些标注不仅覆盖了常见的眼睛、鼻子、嘴巴等区域还包括了头发、眼镜、耳环等细节部分为训练高精度的人脸分割模型提供了坚实基础。1. 数据准备与理解在开始模型训练前我们需要对CelebAMask-HQ数据集有深入理解。这个数据集源自CelebA-HQ但增加了更精细的语义分割标注。每个面部组件都有独立的标注文件最终需要将这些标注合并为单通道的掩码图像。数据集中的19个类别及其对应标签值如下类别名称标签值颜色示例(RGB)皮肤1[204, 0, 0]鼻子2[76, 153, 0]眼睛区域3[204, 204, 0]左眼4[51, 51, 255]右眼5[204, 0, 204]左眉毛6[0, 255, 255]右眉毛7[255, 204, 204]左耳8[102, 51, 0]右耳9[255, 0, 0]嘴巴10[102, 204, 0]上嘴唇11[255, 255, 0]下嘴唇12[0, 0, 153]头发13[0, 0, 204]帽子14[255, 51, 153]耳环15[0, 204, 204]项链16[0, 51, 0]衣服17[255, 153, 51]眼镜18[0, 204, 0]注意数据集中的标签值从1开始0通常保留给背景或未标注区域。在实际训练时可能需要根据模型需求调整标签范围。2. 构建PyTorch数据管道高效的数据加载是模型训练的关键。我们需要自定义一个PyTorch Dataset类来处理CelebAMask-HQ数据import os import torch from torch.utils.data import Dataset from PIL import Image import numpy as np class CelebAMaskHQDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.img_names sorted([f for f in os.listdir(img_dir) if f.endswith(.jpg)]) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_path os.path.join(self.img_dir, self.img_names[idx]) mask_path os.path.join(self.mask_dir, self.img_names[idx].replace(.jpg, .png)) image Image.open(img_path).convert(RGB) mask Image.open(mask_path) # 将彩色mask转换为单通道标签 mask np.array(mask) mask self._color_to_label(mask) if self.transform: image self.transform(image) mask torch.from_numpy(mask).long() return image, mask def _color_to_label(self, color_mask): 将RGB彩色mask转换为单通道标签 label_mask np.zeros((color_mask.shape[0], color_mask.shape[1]), dtypenp.uint8) # 预定义的RGB到标签值的映射 color_to_label { (204, 0, 0): 1, # 皮肤 (76, 153, 0): 2, # 鼻子 (204, 204, 0): 3, # 眼睛区域 (51, 51, 255): 4, # 左眼 (204, 0, 204): 5, # 右眼 (0, 255, 255): 6, # 左眉毛 (255, 204, 204): 7, # 右眉毛 (102, 51, 0): 8, # 左耳 (255, 0, 0): 9, # 右耳 (102, 204, 0): 10, # 嘴巴 (255, 255, 0): 11, # 上嘴唇 (0, 0, 153): 12, # 下嘴唇 (0, 0, 204): 13, # 头发 (255, 51, 153): 14, # 帽子 (0, 204, 204): 15, # 耳环 (0, 51, 0): 16, # 项链 (255, 153, 51): 17, # 衣服 (0, 204, 0): 18 # 眼镜 } for color, label in color_to_label.items(): label_mask[(color_mask color).all(axis2)] label return label_mask使用这个Dataset类时我们可以轻松创建训练和验证集的数据加载器from torchvision import transforms from torch.utils.data import DataLoader # 定义数据增强 train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 创建数据集实例 train_dataset CelebAMaskHQDataset( img_dirpath/to/train_images, mask_dirpath/to/train_masks, transformtrain_transform ) val_dataset CelebAMaskHQDataset( img_dirpath/to/val_images, mask_dirpath/to/val_masks, transformval_transform ) # 创建数据加载器 train_loader DataLoader( train_dataset, batch_size8, shuffleTrue, num_workers4 ) val_loader DataLoader( val_dataset, batch_size8, shuffleFalse, num_workers4 )3. 模型架构选择与实现对于语义分割任务U-Net和DeepLabV3是两种经典且高效的架构。下面我们实现一个改进版的U-Net专门针对人脸分割任务进行了优化import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): (卷积 [BN] ReLU) * 2 def __init__(self, in_channels, out_channels): super().__init__() self.double_conv nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue), nn.Conv2d(out_channels, out_channels, kernel_size3, padding1), nn.BatchNorm2d(out_channels), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.double_conv(x) class Down(nn.Module): 下采样模块 def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_channels, out_channels) ) def forward(self, x): return self.maxpool_conv(x) class Up(nn.Module): 上采样模块 def __init__(self, in_channels, out_channels, bilinearTrue): super().__init__() if bilinear: self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) else: self.up nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 self.up(x1) # 计算填充以确保尺寸匹配 diffY x2.size()[2] - x1.size()[2] diffX x2.size()[3] - x1.size()[3] x1 F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x torch.cat([x2, x1], dim1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super(OutConv, self).__init__() self.conv nn.Conv2d(in_channels, out_channels, kernel_size1) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, n_channels3, n_classes19, bilinearTrue): super(UNet, self).__init__() self.n_channels n_channels self.n_classes n_classes self.bilinear bilinear self.inc DoubleConv(n_channels, 64) self.down1 Down(64, 128) self.down2 Down(128, 256) self.down3 Down(256, 512) self.down4 Down(512, 1024) self.up1 Up(1024, 512, bilinear) self.up2 Up(512, 256, bilinear) self.up3 Up(256, 128, bilinear) self.up4 Up(128, 64, bilinear) self.outc OutConv(64, n_classes) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x5 self.down4(x4) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) logits self.outc(x) return logits这个U-Net实现有几个针对人脸分割的优化点增加了网络的深度和宽度以处理人脸细粒度分割使用双线性插值上采样减少棋盘效应自动处理特征图尺寸不匹配问题输出通道数设置为19对应CelebAMask-HQ的类别数4. 训练策略与评估训练语义分割模型需要考虑几个关键因素损失函数选择、学习率调度和评估指标。以下是完整的训练流程实现import torch.optim as optim from torch.optim import lr_scheduler from tqdm import tqdm def train_model(model, train_loader, val_loader, device, num_epochs50): # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) # 学习率调度器 scheduler lr_scheduler.ReduceLROnPlateau( optimizer, modemax, factor0.5, patience3, verboseTrue ) best_miou 0.0 for epoch in range(num_epochs): model.train() epoch_loss 0.0 # 训练阶段 for images, masks in tqdm(train_loader, descfEpoch {epoch1}/{num_epochs}): images images.to(device) masks masks.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, masks) loss.backward() optimizer.step() epoch_loss loss.item() # 验证阶段 model.eval() val_loss 0.0 total_pixels 0 correct_pixels 0 iou_sum 0.0 with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.to(device) outputs model(images) loss criterion(outputs, masks) val_loss loss.item() # 计算像素准确率 _, predicted torch.max(outputs.data, 1) total_pixels masks.nelement() correct_pixels (predicted masks).sum().item() # 计算各类别的IoU for cls in range(1, 19): # 跳过背景类 pred_mask (predicted cls) true_mask (masks cls) intersection (pred_mask true_mask).sum().float() union (pred_mask | true_mask).sum().float() if union 0: iou_sum (intersection / union).item() # 计算评估指标 pixel_acc correct_pixels / total_pixels mean_iou iou_sum / 17 # 17个非背景类 print(fEpoch {epoch1}: Train Loss: {epoch_loss/len(train_loader):.4f}, fVal Loss: {val_loss/len(val_loader):.4f}, fPixel Acc: {pixel_acc:.4f}, mIoU: {mean_iou:.4f}) # 更新学习率 scheduler.step(mean_iou) # 保存最佳模型 if mean_iou best_miou: best_miou mean_iou torch.save(model.state_dict(), best_model.pth) print(fNew best model saved with mIoU: {best_miou:.4f}) return model在实际训练中我们采用了以下策略来提升模型性能损失函数选择使用标准的交叉熵损失对于类别不平衡问题可以考虑添加权重优化器配置AdamW优化器配合权重衰减防止过拟合学习率调度基于验证集mIoU动态调整学习率评估指标像素级准确率整体分割准确度平均IoUmIoU各类别交并比的平均值更能反映模型性能提示对于人脸分割任务眼镜、头发等小区域的IoU往往较低。可以单独监控这些关键类别的IoU确保模型在这些区域的性能。5. 结果可视化与模型部署训练完成后我们需要验证模型在实际分割任务中的表现。以下是一个可视化函数可以直观展示模型的分割效果import matplotlib.pyplot as plt import numpy as np def visualize_results(model, dataset, device, num_examples3): model.eval() indices np.random.choice(len(dataset), num_examples, replaceFalse) fig, axes plt.subplots(num_examples, 3, figsize(15, 5*num_examples)) if num_examples 1: axes axes[np.newaxis, :] class_colors [ [0, 0, 0], # 背景 [204, 0, 0], # 皮肤 [76, 153, 0], # 鼻子 [204, 204, 0], # 眼睛区域 [51, 51, 255], # 左眼 [204, 0, 204], # 右眼 [0, 255, 255], # 左眉毛 [255, 204, 204], # 右眉毛 [102, 51, 0], # 左耳 [255, 0, 0], # 右耳 [102, 204, 0], # 嘴巴 [255, 255, 0], # 上嘴唇 [0, 0, 153], # 下嘴唇 [0, 0, 204], # 头发 [255, 51, 153], # 帽子 [0, 204, 204], # 耳环 [0, 51, 0], # 项链 [255, 153, 51], # 衣服 [0, 204, 0] # 眼镜 ] for i, idx in enumerate(indices): image, true_mask dataset[idx] # 预测 with torch.no_grad(): input_tensor image.unsqueeze(0).to(device) output model(input_tensor) pred_mask torch.argmax(output, dim1).squeeze().cpu().numpy() # 反归一化图像 image image.permute(1, 2, 0).numpy() mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) image std * image mean image np.clip(image, 0, 1) # 转换为彩色mask true_color np.zeros((*true_mask.shape, 3), dtypenp.uint8) pred_color np.zeros((*pred_mask.shape, 3), dtypenp.uint8) for cls in range(19): true_color[true_mask cls] class_colors[cls] pred_color[pred_mask cls] class_colors[cls] # 绘制结果 axes[i, 0].imshow(image) axes[i, 0].set_title(Input Image) axes[i, 0].axis(off) axes[i, 1].imshow(true_color) axes[i, 1].set_title(Ground Truth) axes[i, 1].axis(off) axes[i, 2].imshow(pred_color) axes[i, 2].set_title(Prediction) axes[i, 2].axis(off) plt.tight_layout() plt.show()对于模型部署我们可以使用TorchScript将模型导出为独立于Python运行时的格式# 导出模型为TorchScript example_input torch.rand(1, 3, 256, 256).to(device) traced_script_module torch.jit.trace(model, example_input) traced_script_module.save(face_segmentation_model.pt)部署后的模型可以集成到各种应用中如虚拟化妆和美容效果应用人脸特效和滤镜视频会议中的背景替换人脸识别和分析系统在实际项目中我发现头发和眼镜区域的分割最具挑战性特别是当头发颜色与背景相似或眼镜有反光时。通过增加这些类别的损失权重和在训练集中添加更多相关样本可以显著提升模型在这些关键区域的表现。