在图像处理与计算机视觉领域图像分割技术一直扮演着至关重要的角色。无论是医学影像分析、自动驾驶场景理解还是工业质检、广告牌识别等实际应用精准分割图像中的目标区域都是核心需求。本周作业训练营聚焦图像分割专题将系统讲解从传统方法到深度学习的主流技术帮助开发者掌握分割任务的核心原理与实战能力。本文将围绕图像分割技术展开涵盖基础概念、经典算法、深度学习模型如FCN、U-Net、YOLO-Seg及完整实战案例。无论你是刚入门计算机视觉的新手还是希望深化分割技术理解的开发者都能通过本文构建系统知识体系并动手实现可复用的分割代码。1. 图像分割技术概述1.1 什么是图像分割图像分割指将数字图像细分为多个图像子区域像素集合的过程其目标是简化或改变图像的表示形式使得图像更容易理解和分析。简单来说分割任务是为每个像素分配一个类别标签从而将图像划分为具有语义意义的区域。与目标检测绘制边界框不同分割要求像素级的精确定位能提供更细致的物体轮廓信息。例如在医学影像中分割可用于精确勾勒肿瘤边界在自动驾驶中分割能识别道路、车辆、行人的每一个像素。1.2 图像分割的主要类型根据任务粒度图像分割可分为语义分割为每个像素分配类别标签不区分同一类别的不同实例如将所有人像素标记为人实例分割区分同一类别的不同实例如区分图像中的多个人全景分割结合语义分割和实例分割同时标注场景中所有物体的类别和实例1.3 传统分割方法简介在深度学习兴起前传统分割方法已发展多年基于阈值的分割利用像素灰度值的差异通过设定阈值分离前景和背景基于边缘的分割检测图像中的边缘信息通过边缘连接形成区域边界基于区域的分割根据像素的相似性进行区域生长或分裂合并基于图论的分割将图像建模为图结构使用最小割/最大流算法进行分割虽然传统方法在某些场景下仍有效但深度学习方法在精度和泛化能力上表现更优已成为当前主流。2. 深度学习图像分割核心模型2.1 全卷积网络FCNFCN是深度学习图像分割的开创性工作其主要贡献在于将传统CNN中的全连接层替换为卷积层使网络可以接受任意尺寸的输入使用反卷积层转置卷积进行上采样恢复空间分辨率引入跳跃连接融合浅层特征和深层特征改善细节恢复FCN的核心思想是编码器-解码器结构编码器通过卷积和池化提取特征解码器通过上采样恢复分辨率。这种结构成为后续众多分割模型的基础。2.2 U-Net网络架构U-Net最初为医学图像分割设计现已成为语义分割的经典模型对称的U型结构左侧为编码器下采样路径右侧为解码器上采样路径跳跃连接将编码器各级特征与解码器对应级特征连接保留空间信息端到端训练整个网络可进行端到端的训练无需预处理或后处理U-Net在医学图像、工业质检等数据量较小的场景表现优异其简洁有效的架构被广泛借鉴和改进。2.3 YOLO-Seg分割模型YOLO-Seg是基于YOLO目标检测框架扩展的分割模型结合了检测和分割的优势实时性能继承YOLO系列的高效推理速度多任务学习同时完成目标检测、实例分割和语义分割端到端优化统一优化检测和分割任务提升整体性能YOLO-Seg适合需要实时分割的应用场景如自动驾驶、视频监控等。3. 环境准备与工具配置3.1 基础环境要求进行图像分割开发需要以下环境配置Python 3.8主流深度学习框架支持的最新版本PyTorch 1.9 或 TensorFlow 2.5根据项目需求选择深度学习框架CUDA 11.0如使用GPU加速确保与深度学习框架版本兼容常用计算机视觉库OpenCV、Pillow、scikit-image等3.2 核心依赖安装创建conda环境并安装必要依赖# 创建并激活conda环境 conda create -n image-segmentation python3.8 conda activate image-segmentation # 安装PyTorch根据CUDA版本选择 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 安装计算机视觉库 pip install opencv-python pillow scikit-image matplotlib numpy # 安装分割相关库 pip install segmentation-models-pytorch albumentations3.3 开发工具配置推荐使用Jupyter Notebook进行实验和调试PyCharm或VS Code进行项目开发# 检查环境是否配置成功 import torch import torchvision import cv2 import numpy as np print(fPyTorch版本: {torch.__version__}) print(fCUDA可用: {torch.cuda.is_available()}) print(fOpenCV版本: {cv2.__version__})4. 基于U-Net的医学图像分割实战4.1 数据集准备与预处理使用ISIC2018皮肤病变分割数据集作为示例import os import numpy as np from PIL import Image import torch from torch.utils.data import Dataset, DataLoader import albumentations as A from albumentations.pytorch import ToTensorV2 class MedicalSegmentationDataset(Dataset): def __init__(self, image_dir, mask_dir, transformNone): 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, idx): img_name self.images[idx] img_path os.path.join(self.image_dir, img_name) mask_path os.path.join(self.mask_dir, img_name.replace(.jpg, _segmentation.png)) image np.array(Image.open(img_path).convert(RGB)) mask np.array(Image.open(mask_path).convert(L), dtypenp.float32) mask[mask 255.0] 1.0 if self.transform is not None: augmentations self.transform(imageimage, maskmask) image augmentations[image] mask augmentations[mask] return image, mask # 数据增强配置 train_transform A.Compose([ A.Resize(256, 256), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) val_transform A.Compose([ A.Resize(256, 256), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])4.2 U-Net模型实现下面是完整的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) self.conv DoubleConv(in_channels, out_channels) else: self.up nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size2, stride2) self.conv DoubleConv(in_channels, out_channels) def forward(self, x1, x2): x1 self.up(x1) # 输入是CHW格式 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_channels, n_classes, 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) factor 2 if bilinear else 1 self.down4 Down(512, 1024 // factor) self.up1 Up(1024, 512 // factor, bilinear) self.up2 Up(512, 256 // factor, bilinear) self.up3 Up(256, 128 // factor, 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 logits4.3 训练流程实现完整的训练脚本包含损失函数、优化器和评估指标import torch.optim as optim from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm def train_model(model, train_loader, val_loader, device, epochs100): # 初始化优化器和损失函数 optimizer optim.Adam(model.parameters(), lr1e-4) criterion nn.BCEWithLogitsLoss() # 二分类分割任务 scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, min, patience5) writer SummaryWriter(runs/medical_segmentation) best_val_loss float(inf) for epoch in range(epochs): model.train() train_loss 0.0 # 训练阶段 for images, masks in tqdm(train_loader, descfEpoch {epoch1}/{epochs}): images images.to(device) masks masks.unsqueeze(1).to(device) # 增加通道维度 optimizer.zero_grad() outputs model(images) loss criterion(outputs, masks) loss.backward() optimizer.step() train_loss loss.item() # 验证阶段 model.eval() val_loss 0.0 with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.unsqueeze(1).to(device) outputs model(images) loss criterion(outputs, masks) val_loss loss.item() avg_train_loss train_loss / len(train_loader) avg_val_loss val_loss / len(val_loader) # 学习率调整 scheduler.step(avg_val_loss) # 记录到tensorboard writer.add_scalar(Loss/train, avg_train_loss, epoch) writer.add_scalar(Loss/val, avg_val_loss, epoch) print(fEpoch {epoch1}: Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}) # 保存最佳模型 if avg_val_loss best_val_loss: best_val_loss avg_val_loss torch.save(model.state_dict(), best_model.pth) writer.close() return model # 初始化模型和训练 device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(n_channels3, n_classes1).to(device) # 假设已准备好数据加载器 # train_loader DataLoader(train_dataset, batch_size8, shuffleTrue) # val_loader DataLoader(val_dataset, batch_size8, shuffleFalse) # 开始训练 # trained_model train_model(model, train_loader, val_loader, device)4.4 模型评估与可视化训练完成后需要对模型性能进行评估import matplotlib.pyplot as plt from sklearn.metrics import jaccard_score, f1_score def evaluate_model(model, test_loader, device): model.eval() iou_scores [] f1_scores [] with torch.no_grad(): for images, masks in test_loader: images images.to(device) masks masks.cpu().numpy() outputs model(images) preds torch.sigmoid(outputs).cpu().numpy() preds (preds 0.5).astype(np.uint8) # 计算评估指标 for i in range(len(masks)): iou jaccard_score(masks[i].flatten(), preds[i].flatten()) f1 f1_score(masks[i].flatten(), preds[i].flatten()) iou_scores.append(iou) f1_scores.append(f1) return np.mean(iou_scores), np.mean(f1_scores) def visualize_results(model, test_loader, device, num_samples5): model.eval() fig, axes plt.subplots(num_samples, 3, figsize(15, 5*num_samples)) with torch.no_grad(): for i, (images, masks) in enumerate(test_loader): if i num_samples: break images images.to(device) outputs model(images) preds torch.sigmoid(outputs).cpu().numpy() preds (preds 0.5).astype(np.uint8) # 显示原图、真实掩码和预测结果 axes[i, 0].imshow(images[0].cpu().permute(1, 2, 0)) axes[i, 0].set_title(Original Image) axes[i, 0].axis(off) axes[i, 1].imshow(masks[0], cmapgray) axes[i, 1].set_title(Ground Truth) axes[i, 1].axis(off) axes[i, 2].imshow(preds[0][0], cmapgray) axes[i, 2].set_title(Prediction) axes[i, 2].axis(off) plt.tight_layout() plt.show() # 评估模型性能 # mean_iou, mean_f1 evaluate_model(model, test_loader, device) # print(fMean IoU: {mean_iou:.4f}, Mean F1: {mean_f1:.4f}) # 可视化结果 # visualize_results(model, test_loader, device)5. 基于FCN的街景分割实战5.1 FCN模型实现使用预训练的VGG16作为编码器实现FCN-8s模型import torchvision.models as models class FCN8s(nn.Module): def __init__(self, num_classes): super(FCN8s, self).__init__() # 使用预训练的VGG16作为特征提取器 vgg16 models.vgg16(pretrainedTrue) features list(vgg16.features.children()) # 编码器部分 self.features_block1 nn.Sequential(*features[:5]) # 池化1后 self.features_block2 nn.Sequential(*features[5:10]) # 池化2后 self.features_block3 nn.Sequential(*features[10:17]) # 池化3后 self.features_block4 nn.Sequential(*features[17:24]) # 池化4后 self.features_block5 nn.Sequential(*features[24:]) # 池化5后 # 分类器调整为卷积 self.classifier nn.Sequential( nn.Conv2d(512, 4096, 7), nn.ReLU(inplaceTrue), nn.Dropout2d(), nn.Conv2d(4096, 4096, 1), nn.ReLU(inplaceTrue), nn.Dropout2d(), nn.Conv2d(4096, num_classes, 1) ) # 上采样和跳跃连接 self.upscore2 nn.ConvTranspose2d(num_classes, num_classes, 4, stride2, biasFalse) self.upscore8 nn.ConvTranspose2d(num_classes, num_classes, 16, stride8, biasFalse) self.upscore_pool4 nn.ConvTranspose2d(num_classes, num_classes, 4, stride2, biasFalse) def forward(self, x): # 前向传播 pool1 self.features_block1(x) pool2 self.features_block2(pool1) pool3 self.features_block3(pool2) pool4 self.features_block4(pool3) pool5 self.features_block5(pool4) # 主要分支 score5 self.classifier(pool5) upscore5 self.upscore2(score5) # 跳跃连接1 score4 self.classifier(pool4) score4 score4[:, :, 5:5upscore5.size(2), 5:5upscore5.size(3)] fuse4 upscore5 score4 upscore4 self.upscore_pool4(fuse4) # 跳跃连接2 score3 self.classifier(pool3) score3 score3[:, :, 9:9upscore4.size(2), 9:9upscore4.size(3)] fuse3 upscore4 score3 # 最终上采样 out self.upscore8(fuse3) out out[:, :, 31:31x.size(2), 31:31x.size(3)].contiguous() return out5.2 街景数据集处理使用Cityscapes数据集进行街景分割class CityscapesDataset(Dataset): def __init__(self, root_dir, splittrain, transformNone, target_transformNone): self.root_dir root_dir self.split split self.transform transform self.target_transform target_transform self.images [] self.targets [] # Cityscapes特定的类别定义和颜色映射 self.classes [road, sidewalk, building, wall, fence, pole, traffic light, traffic sign, vegetation, terrain, sky, person, rider, car, truck, bus, train, motorcycle, bicycle] self.colormap self.create_cityscapes_colormap() # 加载图像和标注路径 image_dir os.path.join(root_dir, leftImg8bit, split) target_dir os.path.join(root_dir, gtFine, split) for city in os.listdir(image_dir): city_image_dir os.path.join(image_dir, city) city_target_dir os.path.join(target_dir, city) for img_name in os.listdir(city_image_dir): if img_name.endswith(_leftImg8bit.png): base_name img_name.replace(_leftImg8bit.png, ) target_name base_name _gtFine_labelIds.png self.images.append(os.path.join(city_image_dir, img_name)) self.targets.append(os.path.join(city_target_dir, target_name)) def create_cityscapes_colormap(self): # Cityscapes官方颜色映射 return { 0: (0, 0, 0), # 未标注 1: (70, 70, 70), # 道路 2: (190, 153, 153), # 人行道 # ... 其他类别颜色映射 } def __len__(self): return len(self.images) def __getitem__(self, idx): image Image.open(self.images[idx]).convert(RGB) target Image.open(self.targets[idx]) if self.transform: image self.transform(image) if self.target_transform: target self.target_transform(target) return image, target6. 图像分割常见问题与解决方案6.1 数据不平衡问题医学图像分割中经常遇到正负样本极度不平衡的情况class DiceLoss(nn.Module): Dice系数损失函数适用于不平衡数据 def __init__(self, smooth1e-6): super(DiceLoss, self).__init__() self.smooth smooth def forward(self, predictions, targets): predictions torch.sigmoid(predictions) predictions predictions.view(-1) targets targets.view(-1) intersection (predictions * targets).sum() dice (2. * intersection self.smooth) / (predictions.sum() targets.sum() self.smooth) return 1 - dice class FocalLoss(nn.Module): Focal Loss解决类别不平衡 def __init__(self, alpha0.8, gamma2): super(FocalLoss, self).__init__() self.alpha alpha self.gamma gamma self.bce nn.BCEWithLogitsLoss(reductionnone) def forward(self, inputs, targets): bce_loss self.bce(inputs, targets) pt torch.exp(-bce_loss) focal_loss self.alpha * (1-pt)**self.gamma * bce_loss return focal_loss.mean() # 组合损失函数 class CombinedLoss(nn.Module): def __init__(self, dice_weight0.5, focal_weight0.5): super(CombinedLoss, self).__init__() self.dice_loss DiceLoss() self.focal_loss FocalLoss() self.dice_weight dice_weight self.focal_weight focal_weight def forward(self, inputs, targets): dice self.dice_loss(inputs, targets) focal self.focal_loss(inputs, targets) return self.dice_weight * dice self.focal_weight * focal6.2 模型过拟合解决方案针对分割模型容易过拟合的问题def create_advanced_augmentations(): 创建强大的数据增强策略 return A.Compose([ A.Resize(256, 256), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.ShiftScaleRotate(shift_limit0.1, scale_limit0.2, rotate_limit45, p0.5), A.RandomBrightnessContrast(p0.3), A.GaussNoise(var_limit(10.0, 50.0), p0.3), A.ElasticTransform(alpha1, sigma50, alpha_affine50, p0.3), A.GridDistortion(p0.3), A.CoarseDropout(max_holes8, max_height32, max_width32, p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) def add_regularization(model, weight_decay1e-4): 为模型添加L2正则化 regularization_loss 0 for param in model.parameters(): regularization_loss torch.norm(param, 2) return weight_decay * regularization_loss6.3 内存优化技巧处理高分辨率图像时的内存优化class MemoryEfficientUNet(UNet): 内存优化的U-Net变体 def __init__(self, n_channels, n_classes, bilinearTrue): super().__init__(n_channels, n_classes, bilinear) def forward(self, x): # 使用梯度检查点减少内存使用 from torch.utils.checkpoint import checkpoint x1 checkpoint(self.inc, x) x2 checkpoint(self.down1, x1) x3 checkpoint(self.down2, x2) x4 checkpoint(self.down3, x3) x5 checkpoint(self.down4, x4) x checkpoint(self.up1, x5, x4) x checkpoint(self.up2, x, x3) x checkpoint(self.up3, x, x2) x checkpoint(self.up4, x, x1) logits self.outc(x) return logits def create_memory_efficient_dataloader(dataset, batch_size4, num_workers2): 创建内存高效的数据加载器 return DataLoader( dataset, batch_sizebatch_size, num_workersnum_workers, pin_memoryTrue if torch.cuda.is_available() else False, persistent_workersTrue if num_workers 0 else False )7. 图像分割评估指标详解7.1 常用评估指标图像分割任务需要多种评估指标全面衡量模型性能import numpy as np from sklearn.metrics import confusion_matrix class SegmentationMetrics: def __init__(self, num_classes): self.num_classes num_classes self.confusion_matrix np.zeros((num_classes, num_classes)) def update(self, pred, target): 更新混淆矩阵 pred pred.flatten() target target.flatten() cm confusion_matrix(target, pred, labelsrange(self.num_classes)) self.confusion_matrix cm def iou(self): 计算每个类别的IoU intersection np.diag(self.confusion_matrix) union self.confusion_matrix.sum(axis1) self.confusion_matrix.sum(axis0) - intersection iou intersection / (union 1e-8) return iou def miou(self): 计算平均IoU iou self.iou() return np.nanmean(iou) def dice(self): 计算Dice系数 intersection np.diag(self.confusion_matrix) sum_pred self.confusion_matrix.sum(axis0) sum_target self.confusion_matrix.sum(axis1) dice (2 * intersection) / (sum_pred sum_target 1e-8) return dice def accuracy(self): 计算准确率 return np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum() def class_accuracy(self): 计算每个类别的准确率 return np.diag(self.confusion_matrix) / (self.confusion_matrix.sum(axis1) 1e-8) # 使用示例 metrics SegmentationMetrics(num_classes2) # 在验证循环中更新指标 # metrics.update(predictions, targets) # print(fmIoU: {metrics.miou():.4f})7.2 可视化分析工具创建全面的可视化分析def plot_training_history(train_losses, val_losses, train_ious, val_ious): 绘制训练历史曲线 fig, (ax1, ax2) plt.subplots(1, 2, figsize(15, 5)) # 损失曲线 ax1.plot(train_losses, labelTraining Loss) ax1.plot(val_losses, labelValidation Loss) ax1.set_title(Training and Validation Loss) ax1.set_xlabel(Epoch) ax1.set_ylabel(Loss) ax1.legend() ax1.grid(True) # IoU曲线 ax2.plot(train_ious, labelTraining IoU) ax2.plot(val_ious, labelValidation IoU) ax2.set_title(Training and Validation IoU) ax2.set_xlabel(Epoch) ax2.set_ylabel(IoU) ax2.legend() ax2.grid(True) plt.tight_layout() plt.show() def create_confusion_matrix_heatmap(metrics, class_names): 创建混淆矩阵热力图 cm metrics.confusion_matrix cm_normalized cm.astype(float) / cm.sum(axis1)[:, np.newaxis] fig, ax plt.subplots(figsize(10, 8)) im ax.imshow(cm_normalized, interpolationnearest, cmapplt.cm.Blues) ax.figure.colorbar(im, axax) # 设置刻度标签 ax.set(xticksnp.arange(cm.shape[1]), yticksnp.arange(cm.shape[0]), xticklabelsclass_names, yticklabelsclass_names, titleNormalized Confusion Matrix, ylabelTrue Label, xlabelPredicted Label) # 旋转x轴标签 plt.setp(ax.get_xticklabels(), rotation45, haright, rotation_modeanchor) # 在格子中显示数值 fmt .2f thresh cm_normalized.max() / 2. for i in range(cm_normalized.shape[0]): for j in range(cm_normalized.shape[1]): ax.text(j, i, format(cm_normalized[i, j], fmt), hacenter, vacenter, colorwhite if cm_normalized[i, j] thresh else black) fig.tight_layout() return fig8. 生产环境部署优化8.1 模型量化与加速将训练好的模型部署到生产环境def quantize_model(model, calibration_loader): 量化模型以减少推理时间 model.eval() # 准备量化配置 model.qconfig torch.quantization.get_default_qconfig(fbgemm) # 准备量化 torch.quantization.prepare(model, inplaceTrue) # 校准 with torch.no_grad(): for images, _ in calibration_loader: model(images) # 转换量化模型 torch.quantization.convert(model, inplaceTrue) return model def optimize_for_inference(model, example_input): 使用TorchScript优化推理 model.eval() # 跟踪模型 traced_model torch.jit.trace(model, example_input) # 保存优化后的模型 torch.jit.save(traced_model, optimized_model.pth) return traced_model class SegmentationAPI: 分割模型推理API def __init__(self, model_path, devicecpu): self.device device self.model self.load_model(model_path) self.preprocess self.create_preprocessing() def load_model(self, model_path): 加载模型 model UNet(n_channels3, n_classes1) model.load_state_dict(torch.load(model_path, map_locationself.device)) model.to(self.device) model.eval() return model def create_preprocessing(self): 创建预处理管道 return A.Compose([ A.Resize(256, 256), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) def predict(self, image): 预测单张图像 with torch.no_grad(): # 预处理 processed self.preprocess(imageimage) input_tensor processed[image].unsqueeze(0).to(self.device) # 推理 output self.model(input_tensor) prediction torch.sigmoid(output).cpu().numpy()[0, 0] # 后处理 mask (prediction 0.5).astype(np.uint8) return mask8.2 批量处理优化处理大量图像时的优化策略class BatchSegmentationProcessor: 批量分割处理器 def __init__(self, model, batch_size8, devicecuda): self.model model self.batch_size batch_size self.device device self