PSPNet 宠物图像语义分割

基于 PyTorch 实现的 PSPNet（Pyramid Scene Parsing Network）语义分割模型，用于宠物图像三分类分割（背景、宠物、边界），采用 ResNet50 空洞卷积骨干网络与金字塔池化模块（PPM），标准输入尺寸 473 x 473。

✨ 功能特性

• PSPNet 经典架构：ResNet50 骨干 + 空洞卷积（dilated convolution）+ 金字塔池化模块，融合多尺度全局上下文信息
• 空洞卷积骨干：layer3 采用 dilation=2，layer4 采用 dilation=4，扩大感受野同时保持特征图分辨率
• 金字塔池化模块（PPM）：四种池化尺度（1x1、2x2、3x3、6x6）捕获不同层级的全局上下文
• 三分类语义分割：针对宠物数据集，分割背景、宠物主体、边界三个类别
• 训练与验证：Adam 优化器、交叉熵损失、自动保存检查点与最佳模型
• 单图推理：输入单张图像，输出原始图与彩色分割掩码对比可视化
• 损坏文件鲁棒处理：数据集加载时自动验证图像与掩码完整性，跳过损坏文件

🛠 技术栈

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

📁 项目结构

PSPNet/
├── models/
│   ├── __init__.py              # 模型包初始化
│   └── pspnet.py                # PSPNet 模型定义（ResNet50 骨干、PPM、解码器）
├── train.py                     # 训练脚本（数据集类、训练循环、验证、检查点保存）
├── predict.py                   # 推理脚本（模型加载、预处理、可视化输出）
├── requirements.txt             # 项目依赖
├── test_image.jpg               # 测试图像示例
├── checkpoints/                 # 模型检查点目录（运行时生成）
│   ├── model_final.pth          # 最终模型权重
│   └── best_model.pth           # 最佳模型权重
└── predictions/                 # 推理输出目录（运行时生成）

🚀 快速开始

环境要求

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

安装

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

数据集准备

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

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

标注图中像素值经 mask = mask - 1 处理后映射为：

• 类别 0：背景
• 类别 1：宠物（前景）
• 类别 2：边界

训练

python train.py

训练配置：

• 输入尺寸：473 x 473（PSPNet 标准输入）
• 优化器：Adam（lr=0.001）
• 损失函数：CrossEntropyLoss
• 批量大小：4
• 数据划分：训练集 80% / 验证集 20%（随机种子 42）
• 模型检查点保存至 checkpoints/model_final.pth
• 最佳模型额外保存为 best_model.pth

推理

python predict.py

默认读取根目录下的 test_image.jpg 进行推理，结果保存至 predictions/ 目录，包含原始图像与彩色分割掩码的对比图（黑色=背景，红色=宠物，绿色=边界）。

📖 模型架构

PSPNet 整体结构

输入图像 (3 x 473 x 473)
    │
    ▼
ResNet50 骨干网络（空洞卷积）
    ├── layer0: conv1 + bn1 + relu + maxpool
    ├── layer1: 标准残差块 (256 通道, stride=1/4)
    ├── layer2: 标准残差块 (512 通道, stride=1/8)
    ├── layer3: 空洞卷积 dilation=2 (1024 通道, stride=1/8)
    └── layer4: 空洞卷积 dilation=4 (2048 通道, stride=1/8)
    │
    ▼
金字塔池化模块（PPM）
    ├── 分支1: AdaptiveAvgPool(1x1) + 1x1 Conv
    ├── 分支2: AdaptiveAvgPool(2x2) + 1x1 Conv
    ├── 分支3: AdaptiveAvgPool(3x3) + 1x1 Conv
    ├── 分支4: AdaptiveAvgPool(6x6) + 1x1 Conv
    └── 原始特征拼接 (2048 x 5 = 10240 通道)
    └── 1x1 Conv 瓶颈层降维至 512 通道 + ReLU
    │
    ▼
Dropout2d(0.3)
    │
    ▼
解码器（三级上采样）
    ├── up_1: Conv(512->256) + BN + ReLU + Dropout(0.1) + 上采样 2x
    ├── up_2: Conv(256->64)  + BN + ReLU + Dropout(0.1) + 上采样 2x
    ├── up_3: Conv(64->64)   + BN + ReLU + Dropout(0.1) + 上采样 2x
    │
    ▼
输出层: Conv(64->3, kernel_size=1) -> 分割图 (3 x 473 x 473)

关键组件

• 空洞卷积（Dilated Convolution）：在 ResNet50 的 layer3 和 layer4 中引入膨胀卷积，layer3 dilation=2、layer4 dilation=4，在不增加参数的前提下扩大感受野，保持输出特征图分辨率为输入的 1/8
• 金字塔池化模块（PPM）：通过四种不同尺度的自适应平均池化（1x1、2x2、3x3、6x6）捕获全局上下文信息，经 1x1 卷积降维后上采样至原始尺寸，与原始特征拼接，有效融合不同层级的场景信息
• 解码器：三级卷积 + 双线性上采样，逐步恢复空间分辨率至原始输入尺寸

📊 效果展示

推理输出示例：

• 左栏：原始输入图像
• 右栏：彩色分割掩码（黑色=背景，红色=宠物，绿色=边界）

仓库根目录附带 test_image.jpg 可直接用于推理测试。

🤝 贡献指南

欢迎提交 Issue 和 Pull Request！

📄 许可证

未指定

👤 作者

GYOUNG - GitHub