概率 U-Net

+ **更新**
+ 现已提供改进后的模型（层次概率 U-Net）及 LIDC 裁剪数据集。详情见下文。

重新实现论文《用于模糊图像分割的概率 U-Net》中描述的模型（NeurIPS 2018 论文）。

该模型也在 NeurIPS 上进行了亮点展示，一篇内容相似的简短视频可在此处观看：这里（4 分钟）。

概率 U-Net 的架构如下所示：子图 a) 展示采样过程，b) 展示训练设置：

以下是基于 CityScapes 数据集中保留的验证集图像生成的样本：

在虚拟环境中设置软件包

git clone https://github.com/SimonKohl/probabilistic_unet.git .
cd prob_unet/
virtualenv -p python3 venv
source venv/bin/activate
pip3 install -e .

安装用于数据增强的批处理生成器

cd ..
git clone https://github.com/MIC-DKFZ/batchgenerators
cd batchgenerators
pip3 install nilearn scikit-image nibabel
pip3 install -e .
cd prob_unet

下载并预处理 CityScapes 数据集

1. 在 CityScapes 官网注册账号：https://www.cityscapes-dataset.com/
2. 登录后，下载训练、验证和测试的标注文件及图像：
   • 标注文件：gtFine_trainvaltest.zip（241MB）
   • 图像：leftImg8bit_trainvaltest.zip（11GB）
3. 解压数据（解压 _trainvaltest.zip），并在 preprocessing_config.py 中调整 raw_data_dir（解压后文件的完整路径）和 out_dir（期望输出目录的完整路径）。
4. 使用双线性插值将数据缩放到 256 x 512 分辨率，并以 NumPy 数组格式保存，运行以下命令：

cd cityscapes
python3 preprocessing.py
cd ..

训练

[如果您只想使用预训练模型，请跳至评估部分。]
修改 scripts/prob_unet_config.py 中的 data_dir 和 exp_dir，然后执行：

cd training
python3 train_prob_unet.py --config prob_unet_config.py

规格说明

加载您自己训练好的模型或使用预训练模型。一组预训练权重可从 zenodo.org 下载（187MB）。下载后，通过以下命令解压文件： tar -xvzf pretrained_weights.tar.gz，例如解压到 /model 目录。无论使用您自己的模型还是预训练模型，都需在 evaluation/cityscapes_eval_config.py 中修改 data_dir 和 exp_dir，使其与您的路径一致。

首先生成样本（默认为 500 张验证图像中的每张生成 16 个分割样本）：

cd ../evaluation
python3 eval_cityscapes.py --write_samples

随后进行评估（多线程进行，因此速度较快）：

python3 eval_cityscapes.py --eval_samples

评估会生成一个包含结果的字典。可通过启动 IPython 笔记本进行可视化：

jupyter notebook evaluation_plots.ipynb

使用上述笔记本对预训练模型进行评估后得到的结果如下：

测试

评估指标已纳入测试覆盖范围。运行测试如下：

cd ../tests/evaluation
python3 -m pytest eval_tests.py

与原始工作的差异

本仓库中的代码并未用于原始论文，且存在一些细微改动：

• 在单块 GPU（Titan Xp）上进行训练，而非分布式训练；当前实现不支持分布式训练。
• 模型中的下采样操作采用平均池化而非双线性插值。
• 卷积核数量在第三尺度后保持不变，而非严格地每经过一个尺度就翻倍（以减少内存占用）。
• HeNormal 权重初始化的效果优于正交权重初始化。

如何引用此代码

请引用原始出版物：

@article{kohl2018probabilistic,
  title={A Probabilistic U-Net for Segmentation of Ambiguous Images},
  author={Kohl, Simon AA and Romera-Paredes, Bernardino and Meyer, Clemens and De Fauw, Jeffrey and Ledsam, Joseph R and Maier-Hein, Klaus H and Eslami, SM and Rezende, Danilo Jimenez and Ronneberger, Olaf},
  journal={arXiv preprint arXiv:1806.05034},
  year={2018}
}

许可证

代码根据 Apache License Version 2.0 发布。

更新：层次概率 U-Net + LIDC 裁剪数据集

我们在 2019 年的 Medical Imaging meets Neurips 研讨会上发表了一种改进后的模型——层次概率 U-Net。

论文可在 arXiv 上获取：《用于建模多尺度模糊性的层次概率 U-Net》，2019 年 5 月。

模型代码可从 DeepMind 的 GitHub 仓库免费获取，链接如下：代码链接。

LIDC 数据可从 Google Cloud Storage 下载为尺寸为 180 x 180 的 PNG 文件，链接如下：数据链接。

预训练模型可直接应用于该数据集，只需使用以下 Google Colab：。

Probabilistic U-Net 快速上手指南

Probabilistic U-Net 是 NeurIPS 2018 亮点论文《A Probabilistic U-Net for Segmentation of Ambiguous Images》的复现实现。该模型专为处理具有模糊性或多义性的图像分割任务设计，能够生成多种合理的分割结果。

环境准备

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: Python 3.x
• 硬件要求: 支持 CUDA 的 NVIDIA GPU（原论文使用 Titan Xp 单卡训练）
• 前置依赖:
  • virtualenv (用于创建虚拟环境)
  • git
  • pip3

安装步骤

1. 克隆代码并配置虚拟环境

git clone https://github.com/SimonKohl/probabilistic_unet.git .
cd prob_unet/
virtualenv -p python3 venv
source venv/bin/activate
pip3 install -e .

2. 安装数据增强工具 (batch-generators)

该项目依赖 MIC-DKFZ 的 batch-generators 库进行数据增强。

cd ..
git clone https://github.com/MIC-DKFZ/batchgenerators
cd batchgenerators
# 安装必要的图像处理依赖
pip3 install nilearn scikit-image nibabel
pip3 install -e .
cd prob_unet

提示: 国内用户若遇到 pip 下载缓慢，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 参数加速安装。

3. 数据准备 (以 Cityscapes 数据集为例)

1. 访问 Cityscapes 官网 注册账号。
2. 下载以下两个数据包：
   • 标注文件: gtFine_trainvaltest.zip (241MB)
   • 图像文件: leftImg8bit_trainvaltest.zip (11GB)
3. 解压文件，并修改 preprocessing_config.py 中的路径配置：
   • raw_data_dir: 指向解压后的文件根目录
   • out_dir: 指向期望的输出目录
4. 执行预处理脚本（将数据双线性缩放至 256x512 并保存为 numpy 数组）：

cd cityscapes
python3 preprocessing.py
cd ..

基本使用

场景一：训练模型

如果您希望从头训练模型，请修改配置文件 scripts/prob_unet_config.py 中的 data_dir (数据目录) 和 exp_dir (实验输出目录)，然后运行：

cd training
python3 train_prob_unet.py --config prob_unet_config.py

场景二：使用预训练模型进行评估

如果您只想快速体验或评估效果，可以使用官方提供的预训练权重。

1. 下载权重: 从 Zenodo 下载 pretrained_weights.tar.gz (187MB)。
2. 解压权重:

   tar -xvzf pretrained_weights.tar.gz
   # 建议解压到 /model 目录
   

3. 配置路径: 修改 evaluation/cityscapes_eval_config.py 中的 data_dir 和 exp_dir 以匹配您的实际路径。
4. 生成分割样本 (默认为 500 张验证图各生成 16 个分割样本):

   cd evaluation
   python3 eval_cityscapes.py --write_samples
   

5. 执行评估 (多线程运行，速度较快):

   python3 eval_cityscapes.py --eval_samples
   

6. 可视化结果: 评估结果将保存在字典中，可通过 Jupyter Notebook 查看可视化图表：

   jupyter notebook evaluation_plots.ipynb
   

---

注：本项目代码基于 Apache License Version 2.0 开源。如需引用，请参考原论文 Kohl et al., arXiv:1806.05034 (2018)。