# 导入核心依赖库 import torch # PyTorch核心库用于张量计算和模型构建 import torch.nn as nn # 神经网络层模块 import torch.nn.functional as F # 神经网络函数如激活、损失计算 from PIL import Image # Python Imaging Library用于读取/处理图片 import numpy as np # 数值计算库处理数组 import matplotlib.pyplot as plt # 可视化库展示图片和分割结果 import cv2 # OpenCV库用于保存掩码图片 # 1. 定义U-Net核心网络结构带详细注释 class DoubleConv(nn.Module): U-Net基础卷积块连续两次卷积BNReLU保证特征提取充分 作用每个编码器/解码器阶段的核心特征提取单元 def __init__(self, in_channels, out_channels): super().__init__() # 继承nn.Module的初始化 # 构建卷积序列Conv2d→BN→ReLU 执行两次 self.double_conv nn.Sequential( # 第一次卷积same paddingpadding1保证卷积后尺寸不变 # in_channels输入通道数out_channels输出通道数kernel_size33×3卷积核特征提取核心 nn.Conv2d(in_channels, out_channels, kernel_size3, padding1), # 批量归一化BN加速训练收敛防止梯度消失/爆炸 nn.BatchNorm2d(out_channels), # ReLU激活函数引入非线性inplaceTrue节省内存 nn.ReLU(inplaceTrue), # 第二次卷积进一步提取特征通道数保持out_channels不变 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): U-Net编码器下采样块最大池化尺寸减半 DoubleConv特征翻倍 作用缩小特征图尺寸增加通道数提取全局特征 def __init__(self, in_channels, out_channels): super().__init__() self.maxpool_conv nn.Sequential( # 最大池化步长2特征图尺寸减半如128×128→64×64 nn.MaxPool2d(2), # 池化后接DoubleConv通道数从in_channels→out_channels如64→128 DoubleConv(in_channels, out_channels) ) def forward(self, x): 前向传播输入→池化→卷积→输出 return self.maxpool_conv(x) class Up(nn.Module): U-Net解码器上采样块转置卷积尺寸翻倍 跳跃连接拼接 DoubleConv 作用恢复特征图尺寸融合编码器的局部特征和解码器的全局特征 def __init__(self, in_channels, out_channels): super().__init__() # 转置卷积反卷积上采样尺寸翻倍通道数减半如512→256 # kernel_size2/stride2保证尺寸刚好翻倍如64×64→128×128 self.up nn.ConvTranspose2d(in_channels, in_channels//2, kernel_size2, stride2) # 拼接后接DoubleConv融合跳跃连接的特征通道数最终为out_channels self.conv DoubleConv(in_channels, out_channels) def forward(self, x1, x2): 前向传播 x1解码器上一层输出小尺寸如64×64 x2编码器对应层输出大尺寸如128×128用于跳跃连接 x1 self.up(x1) # x1上采样尺寸翻倍和x2一致 # 跳跃连接核心拼接x2编码器局部特征和x1解码器全局特征 # dim1按通道维度拼接如x2256通道x1256通道→拼接后512通道 x torch.cat([x2, x1], dim1) return self.conv(x) # 拼接后卷积融合特征 class OutConv(nn.Module): U-Net输出层1×1卷积 作用不改变尺寸仅调整通道数为类别数如64通道→2类 def __init__(self, in_channels, out_channels): super(OutConv, self).__init__() # 1×1卷积核尺寸1仅改变通道数无尺寸变化 self.conv nn.Conv2d(in_channels, out_channels, kernel_size1) def forward(self, x): return self.conv(x) class UNet(nn.Module): 完整U-Net模型编码器→解码器→输出层 def __init__(self, n_channels, n_classes): super(UNet, self).__init__() self.n_channels n_channels # 输入图片通道数RGB3灰度图1 self.n_classes n_classes # 分割类别数二分类2多分类按需设置 # 编码器下采样 self.inc DoubleConv(n_channels, 64) # 输入层n_channels→64通道尺寸不变 self.down1 Down(64, 128) # 下采样164→128通道尺寸减半 self.down2 Down(128, 256) # 下采样2128→256通道尺寸再减半 self.down3 Down(256, 512) # 下采样3256→512通道尺寸再减半 # 解码器上采样 self.up1 Up(512, 256) # 上采样1512→256通道尺寸翻倍 self.up2 Up(256, 128) # 上采样2256→128通道尺寸翻倍 self.up3 Up(128, 64) # 上采样3128→64通道尺寸翻倍回到原始尺寸 # 输出层 self.outc OutConv(64, n_classes) # 64通道→类别数输出分割预测 def forward(self, x): U-Net完整前向传播编码器→保存跳跃连接→解码器→输出 # 编码器前向保存各层输出用于跳跃连接 x1 self.inc(x) # 原始尺寸64通道 x2 self.down1(x1) # 尺寸减半128通道 x3 self.down2(x2) # 尺寸再减半256通道 x4 self.down3(x3) # 尺寸再减半512通道 # 解码器前向拼接跳跃连接 x self.up1(x4, x3)# 上采样拼接x3256通道尺寸翻倍 x self.up2(x, x2) # 上采样拼接x2128通道尺寸翻倍 x self.up3(x, x1) # 上采样拼接x164通道尺寸翻倍回到原始尺寸 # 输出层预测每个像素的类别概率shape(batch, n_classes, H, W) logits self.outc(x) return logits # 2. 图片预处理/后处理函数推理核心 def preprocess_image(img_path, target_size(128, 128)): 预处理本地图片适配模型输入格式 :param img_path: 本地图片路径如test.jpg :param target_size: 模型输入尺寸需和训练时一致否则报错 :return: img_tensor: 预处理后的张量shape(1, 3, H, W)batch1通道在前 img_np: 原始numpy格式图片用于可视化 # 1. 读取图片PIL打开并转为RGB模式避免灰度图/透明通道问题 img Image.open(img_path).convert(RGB) # 2. 调整尺寸模型输入要求固定尺寸resize统一为target_size img img.resize(target_size) # 3. 格式转换PIL→numpy数组H, W, C→PyTorch张量C, H, W img_np np.array(img).astype(np.float32) # PIL图片是(H,W,C)转float32避免精度损失 img_tensor torch.from_numpy(img_np).permute(2, 0, 1) # 维度转换(H,W,C)→(C,H,W)PyTorch要求通道在前 # 4. 归一化像素值从0-255→0-1和训练时的预处理一致否则模型预测异常 img_tensor img_tensor / 255.0 # 5. 增加batch维度模型输入要求(batch_size, C, H, W)单张图片batch_size1 img_tensor img_tensor.unsqueeze(0) return img_tensor, img_np # 返回张量推理用和原始numpy图可视化用 def postprocess_pred(pred_tensor): 后处理模型预测结果转为可视化的分割掩码 :param pred_tensor: 模型输出shape(1, n_classes, H, W) :return: pred_mask: 分割掩码shape(H, W)每个像素值为类别编号 # 1. 取每个像素概率最大的类别dim1表示按通道维度类别维度取最大值 pred_mask torch.argmax(pred_tensor, dim1) # 2. 去除batch维度从(1, H, W)→(H, W)转numpy数组可视化/保存用 pred_mask pred_mask.squeeze(0).cpu().numpy() return pred_mask # 3. 完整推理流程函数 def infer_local_image(img_path, n_classes2, target_size(128, 128)): 加载模型→预处理图片→推理→后处理→可视化保存结果 :param img_path: 本地图片路径 :param n_classes: 分割类别数二分类2 :param target_size: 模型输入尺寸 :return: pred_mask: 分割掩码数组 # 1. 设备选择优先用GPUCUDA无则用CPU推理速度GPU远快于CPU device torch.device(cuda if torch.cuda.is_available() else cpu) print(f使用设备: {device}) # 2. 初始化模型并加载可选加载训练权重 # 实例化U-Net输入通道3RGB类别数n_classes model UNet(n_channels3, n_classesn_classes).to(device) # 模型移到指定设备CPU/GPU # 【关键注释】如果有训练好的权重取消下面注释加载否则用随机初始化的模型预测结果随机 # model.load_state_dict(torch.load(unet_trained.pth, map_locationdevice)) # 加载训练权重 model.eval() # 切换到推理模式关闭BN/Dropout的训练行为避免预测结果异常 # 3. 预处理图片转为模型可识别的张量 img_tensor, img_np preprocess_image(img_path, target_size) img_tensor img_tensor.to(device) # 图片张量移到指定设备 # 4. 模型推理核心步骤 with torch.no_grad(): # 关闭梯度计算节省内存加快推理速度推理无需反向传播 pred_tensor model(img_tensor) # 输入图片→模型→预测结果 # 5. 后处理预测结果转为分割掩码 pred_mask postprocess_pred(pred_tensor) # 6. 可视化结果原图 分割掩码对比展示 plt.figure(figsize(10, 5)) # 设置画布大小 # 子图1原始图片预处理后的尺寸 plt.subplot(1, 2, 1) # 1行2列第1个子图 plt.imshow(img_np.astype(np.uint8)) # 转回0-255的uint8格式才能正常显示 plt.title(Original Image) # 标题 plt.axis(off) # 关闭坐标轴更美观 # 子图2分割掩码jet色板0背景蓝色1目标红色 plt.subplot(1, 2, 2) # 1行2列第2个子图 plt.imshow(pred_mask, cmapjet) # cmap指定色板区分不同类别 plt.title(Segmentation Mask) plt.axis(off) # 保存可视化结果bbox_inchestight去除白边dpi150保证清晰度 plt.savefig(segmentation_result.png, bbox_inchestight, dpi150) plt.show() # 显示图片 # 7. 保存纯分割掩码转为0-255的灰度图方便后续处理如抠图 mask_img (pred_mask * 255).astype(np.uint8) # 类别0→0黑类别1→255白 cv2.imwrite(segmentation_mask.png, mask_img) # OpenCV保存灰度图 print(分割结果已保存) print(- 可视化对比图segmentation_result.png) print(- 纯分割掩码图segmentation_mask.png) return pred_mask # 返回分割掩码供后续分析 # 4. 执行推理主函数 if __name__ __main__: # 配置参数替换为你的本地图片路径 img_path test.jpg # 图片放在代码同目录下否则写绝对路径如C:/Users/xxx/Desktop/test.jpg # 执行推理n_classes2二分类背景目标target_size128×128和模型输入一致 infer_local_image(img_path, n_classes2, target_size(128, 128))