U-Net 宠物图像语义分割

基于 PyTorch 实现的 U-Net 语义分割模型，用于宠物图像三分类分割（背景、宠物、边界），采用经典编码器-解码器架构与跳跃连接，配套 Flask Web 演示界面与 mIoU 评估工具。

✨ 功能特性

• U-Net 经典架构：4 层下采样编码器 + 4 层上采样解码器，通过跳跃连接融合多尺度特征
• 三分类语义分割：针对宠物数据集，分割背景、宠物主体、边界三个类别
• 模型训练：Adam 优化器、交叉熵损失、自动保存模型权重
• 单图推理：输入单张图像，输出原始图与分割掩码对比可视化
• Flask Web 演示：浏览器上传图像，实时返回彩色分割结果
• mIoU 评估：逐类别计算 IoU 与整体 mIoU，输出柱状图与文本报告
• 置信度阈值分析：提供阈值对分割性能影响的分析脚本

🛠 技术栈

类别	技术
语言	Python 3.x
深度学习框架	PyTorch >= 1.7.0
计算机视觉	torchvision >= 0.8.0, OpenCV >= 4.5.0, Pillow >= 8.0.0
科学计算	NumPy >= 1.19.0
可视化	Matplotlib >= 3.3.0
Web 框架	Flask >= 2.0.0
工具	tqdm >= 4.50.0, pycocotools >= 2.0.0

📁 项目结构

U-Net/
├── train.py                       # U-Net 模型定义 + 训练脚本 + 数据集类
├── predict.py                     # 单图推理脚本（加载模型、预处理、可视化）
├── app.py                         # Flask Web 应用（上传图像、实时分割）
├── calculate_miou.py              # mIoU 评估脚本（逐类别 IoU、整体 mIoU）
├── analyze_confidence_threshold.py # 置信度阈值影响分析脚本
├── simple_miou_calculator.py      # 简化版 mIoU 计算工具
├── requirements.txt               # 项目依赖
├── templates/
│   └── index.html                 # Web 演示页面
├── static/
│   ├── uploads/                   # 用户上传图像（运行时生成）
│   └── results/                   # 分割结果图像（运行时生成）
├── predictions/                   # 推理输出目录（运行时生成）
├── miou_results.txt               # mIoU 评估结果文本
├── threshold_impact_results.txt   # 阈值分析结果文本
└── unet_pet.pth                   # 训练好的模型权重（需训练获得）

🚀 快速开始

环境要求

• Python >= 3.7
• CUDA（可选，用于 GPU 加速）

安装

git clone https://github.com/G-YOUNG01/U-Net.git
cd U-Net
pip install -r requirements.txt

数据集准备

本项目使用 Oxford-IIIT Pet Dataset 格式，需按以下目录结构放置数据：

data/
├── images/                  # 原始宠物图像（.jpg）
└── annotations/
    └── trimaps/             # 三值分割标注（.png，像素值 1/2/3）

标注图中像素值映射关系：

• 原始标签 1（前景）→ 类别 0：背景
• 原始标签 2（背景）→ 类别 1：宠物
• 原始标签 3（边界）→ 类别 2：边界

训练

python train.py

训练配置：

• 输入尺寸：256 × 256
• 优化器：Adam（lr=1e-4）
• 损失函数：CrossEntropyLoss
• 批量大小：4
• 模型权重保存为 unet_pet.pth

单图推理

python predict.py

默认读取 test_image.jpg 进行推理，结果保存至 predictions/ 目录，包含原始图像与预测掩码的对比图。

Web 演示

python app.py

启动后访问 http://localhost:5000，在页面上传宠物图像即可实时获得彩色分割结果（黑色=背景，红色=宠物，绿色=边界）。

mIoU 评估

python calculate_miou.py

脚本会加载 unet_pet.pth，在验证集上计算每个类别的 IoU 和整体 mIoU，输出：

• 控制台打印各类别 IoU 与平均 mIoU
• miou_results.png：各类别 mIoU 柱状图
• miou_results.txt：评估结果文本文件

📖 模型架构

U-Net 整体结构

输入图像 (3 × 256 × 256)
    │
    ▼
编码器（下采样路径）
    ├── inc:    DoubleConv(3→64)         ──► x1 (64 × 256 × 256)
    ├── down1:  MaxPool + DoubleConv(64→128)  ──► x2 (128 × 128 × 128)
    ├── down2:  MaxPool + DoubleConv(128→256) ──► x3 (256 × 64 × 64)
    ├── down3:  MaxPool + DoubleConv(256→512) ──► x4 (512 × 32 × 32)
    └── down4:  MaxPool + DoubleConv(512→1024) ──► x5 (1024 × 16 × 16)
    │
    ▼
解码器（上采样路径）
    ├── up1: ConvTranspose2d(1024→512) + cat(x4) + DoubleConv(1024→512)
    ├── up2: ConvTranspose2d(512→256)  + cat(x3) + DoubleConv(512→256)
    ├── up3: ConvTranspose2d(256→128)  + cat(x2) + DoubleConv(256→128)
    ├── up4: ConvTranspose2d(128→64)   + cat(x1) + DoubleConv(128→64)
    │
    ▼
输出层: Conv2d(64→3, kernel_size=1) → 分割图 (3 × 256 × 256)

关键组件

• DoubleConv：两个连续的 3×3 卷积 + BatchNorm + ReLU，是 U-Net 的基本构建块
• 跳跃连接：解码器每一层通过通道拼接（concat）融合编码器对应层级的特征，保留精细的空间信息
• 转置卷积上采样：使用 ConvTranspose2d 实现 2 倍上采样，恢复空间分辨率

📊 效果展示

mIoU 评估结果

类别	IoU
背景	0.7109
宠物	0.7914
边界	0.3848
平均 mIoU	0.6290

注：边界类别 IoU 较低是语义分割任务中的常见现象，因边界像素占比小且标注存在主观性。

推理输出

• 原始图像与预测掩码左右对比图
• Web 演示返回彩色分割叠加图

🤝 贡献指南

欢迎提交 Issue 和 Pull Request！

📄 许可证

未指定

👤 作者

GYOUNG - GitHub