Water Body Segmentation

基于 U-Net++ 的水域图像分割与预测可视化。

本项目是危险水域自动标注研究中的图像分割模块：输入 RGB 图像，输出水体的像素级二值掩码，并将预测结果与人工标注并排展示。仓库覆盖数据加载、模型构建、组合损失、训练验证和预测流程，适合学习与复现水体分割实验。

Water-body segmentation with a ResNet50-backed U-Net++, a combined BCE/Dice/Focal loss, and an end-to-end training and visualization workflow.

技术实现

环节	当前代码实现
分割模型	segmentation_models_pytorch.UnetPlusPlus，ResNet50 编码器
初始化	ImageNet 编码器预训练权重
输入与输出	RGB 图像；单通道 logits，经 sigmoid 和阈值转换为水体掩码
损失函数	0.3 × BCE + 0.5 × Dice + 0.2 × Focal
优化	AdamW、ReduceLROnPlateau 学习率调度
评估	阈值化掩码的 IoU
展示	原图、预测掩码、真实标签三列对照

文件结构

config.py       # 数据路径、设备及超参数
dataset.py      # 图像与标签读取、变换
data_loader.py  # 训练和验证 DataLoader
model.py        # U-Net++ 模型封装
loss.py         # BCE / Dice / Focal 组合损失
trainer.py      # 训练、验证、权重保存
predict.py      # 随机样本预测与可视化

环境准备

建议使用独立 Python 环境，并按显卡与 CUDA 环境安装匹配的 PyTorch。

python -m venv .venv
pip install -r requirements.txt

requirements.txt 列出源码依赖，未锁定原训练环境版本；首次运行使用 ImageNet 初始化时可能下载编码器权重。

数据准备

WaterBodies/
├── Image/       # 训练图像
├── Mask/        # 训练二值掩码
├── val_image/   # 独立验证图像
└── val_mask/    # 独立验证掩码

图像支持 .jpg、.png、.tif；标签优先使用同名 .png，否则尝试与图像同名。水体为白色，背景为黑色。请确保每张图像都有正确配对的标签，并划分互不重叠的训练与验证数据。

训练与预测

1. 在 config.py 设置 dataset_path、训练/验证图像与掩码目录，以及 model_save_dir。原代码保留了开发机器的绝对路径，运行前必须替换。
2. 根据显存调整 batch_size，再运行训练。
3. 确认 best_model_path 指向生成的模型权重后运行预测。

python trainer.py
python predict.py

默认配置为 256 × 256 输入、batch size 6、10 个 epoch、学习率 0.0004、预测阈值 0.5。预测脚本读取 dataset_path/Image 与 dataset_path/Mask，并保存 random_samples_prediction.png；该可视化不是独立测试集指标。

当前版本与验证范围

本仓库发布的是原始分割模块源码，未包含数据集和模型权重，本次整理只进行了 Python 语法检查，未重新训练或验证精度。以下事项需要在正式实验前处理：

• dataset.py 对图像和掩码分别采样随机几何变换，可能导致训练标签错位；应改为同步变换后再训练和报告指标。
• 缺失标签时当前代码生成全黑掩码；训练前应检查配对完整性。
• model.py 的 deep_supervision 是一个标记，当前 forward 没有返回多尺度监督输出。
• 这里实现的是水体分割；像素到经纬度转换、地理信息系统集成和危险等级判定不在此版本中。

因此，仓库不将申请材料中的系统级精度或地理标注能力作为本代码版本的已验证结果。

Acknowledgements

Built with PyTorch, torchvision, and segmentation-models-pytorch. Dependencies retain their respective licenses.