CS60003-hw2_3

U-Net图像分割 - 任务三

项目简介

本项目从零实现U-Net语义分割网络，在Stanford Background Dataset上进行8类语义分割任务，对比三种损失函数（交叉熵损失、Dice Loss、组合损失）的性能表现。

主要特点：

• 完全手写U-Net，无任何预训练权重
• 支持8类语义分割（天空、树木、道路、草地、水面、建筑物、山脉、前景）
• 实现三种损失函数对比
• 支持忽略索引（处理未标注区域-1, -2等）
• 集成SwanLab可视化

项目结构

hw_2_3/
├── config.py                     # 配置文件
├── models.py                     # U-Net模型（手写，无预训练）
├── dataset.py                    # 数据集加载
├── losses.py                     # 三种损失函数
├── metrics.py                    # 评估指标
├── train.py                      # 训练逻辑
├── main.py                       # 主程序
├── visualize.py                  # 可视化工具
├── requirements.txt              # 依赖包
├── StanfordBackgroundDataset/    # 数据集文件夹
│   ├── images/...               # 图片文件夹
│   └── labels/...               # 标签文件夹
└── results/                      # 结果保存目录

数据集

Stanford Background Dataset (SBD)

类别ID	类别名称	原始标签
0	sky (天空)	0
1	tree (树木)	1
2	road (道路)	2
3	grass (草地)	3
4	water (水面)	4
5	building (建筑物)	5
6	mountain (山脉)	6
7	foreground (前景)	>=7

训练配置

参数	值
训练/验证划分	80% / 20%
图像尺寸	256x256
Batch Size	8
Epochs	50
学习率	1e-4
优化器	Adam
损失函数	CE / Dice / Combined

快速开始

# 1. 安装依赖
pip install -r requirements.txt

# 2. 运行训练
python main.py

输出文件

results/
├── models/
│   ├── best_model_ce.pth
│   ├── best_model_dice.pth
│   └── best_model_combined.pth
└── curves/
    ├── loss_curves.png
    ├── miou_curves.png
    ├── accuracy_curves.png
    └── combined_curves.png

依赖包

torch>=1.9.0
torchvision>=0.10.0
numpy>=1.19.0
Pillow>=8.0.0
matplotlib>=3.3.0
tqdm>=4.50.0
swanlab>=0.3.0