PyTorch-3DUNet：基于PyTorch实现的3D U-Net模型用于体积语义分割

在计算机视觉和医学影像分析领域，3D U-Net模型因其在处理体积数据方面的优异表现而备受关注。今天，我们将为大家详细介绍一个基于PyTorch实现的3D U-Net开源项目——PyTorch-3DUNet。该项目为研究人员和开发者提供了一个强大而灵活的工具，用于解决各种体积语义分割问题。

项目概述

PyTorch-3DUNet是由GitHub用户wolny开发并维护的开源项目。该项目在GitHub上已获得1.9k颗星和489次分叉，显示了其在社区中的受欢迎程度和实用价值。项目遵循MIT许可证，为用户提供了充分的使用和修改自由。

PyTorch-3DUNet Logo

主要特性

PyTorch-3DUNet项目实现了多种3D U-Net变体，包括：

标准3D U-Net：基于论文《3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation》实现。
残差3D U-Net：参考《Superhuman Accuracy on the SNEMI3D Connectomics Challenge》论文实现。
带SE模块的残差3D U-Net：在残差U-Net基础上增加了Squeeze-and-Excitation模块，灵感来自《Deep Learning Semantic Segmentation for High-Resolution Medical Volumes》和《Squeeze-and-Excitation Networks》论文。

这些模型不仅可以用于语义分割任务，还支持回归问题（如去噪和学习反卷积）。

输入数据格式

PyTorch-3DUNet支持处理存储在HDF5文件中的输入数据。训练数据应包含两个数据集：raw和label。对于使用PixelWiseCrossEntropyLoss进行训练时，还需要提供weight数据集。数据格式因问题维度和通道数而异：

2D单通道：(1, Y, X)
2D多通道：(C, 1, Y, X)
3D单通道：(Z, Y, X)
3D多通道：(C, Z, Y, X)

安装和使用

环境要求

NVIDIA GPU
CUDA CuDNN

安装方法

最简便的安装方式是通过conda/mamba：

conda install -c conda-forge mamba
mamba create -n pytorch-3dunet -c pytorch -c nvidia -c conda-forge pytorch pytorch-cuda=12.1 pytorch-3dunet
conda activate pytorch-3dunet

安装完成后，可以使用train3dunet命令进行模型训练，使用predict3dunet命令进行预测。

模型训练

通过以下命令启动训练：

train3dunet --config <CONFIG>

其中<CONFIG>是一个YAML配置文件的路径，用于指定训练过程的各个方面。项目提供了多个示例配置文件，如用于3D语义分割的train_config_segmentation.yaml和用于3D回归任务的train_config_regression.yaml。

模型预测

使用以下命令进行预测：

predict3dunet --config <CONFIG>

同样，需要提供一个配置文件来指定预测过程的参数。

数据并行

PyTorch-3DUNet默认支持多GPU训练和预测。如果有多个GPU可用，训练和预测将自动在所有GPU上运行。用户可以通过设置CUDA_VISIBLE_DEVICES环境变量来限制使用的GPU数量。

支持的损失函数

PyTorch-3DUNet支持多种损失函数，适用于不同的任务类型：

语义分割

BCEWithLogitsLoss
DiceLoss
BCEDiceLoss
CrossEntropyLoss
PixelWiseCrossEntropyLoss
WeightedCrossEntropyLoss
GeneralizedDiceLoss

回归

MSELoss
L1Loss
SmoothL1Loss
WeightedSmoothL1Loss

评估指标

项目还提供了多种评估指标来衡量模型性能：

语义分割

MeanIoU（平均交并比）
DiceCoefficient（Dice系数）
BoundaryAveragePrecision（边界平均精度）
AdaptedRandError
AveragePrecision（平均精度）

回归

PSNR（峰值信噪比）
MSE（均方误差）

应用实例

PyTorch-3DUNet项目提供了多个实际应用的示例，包括：

拟南芥侧根光片显微图像的细胞边界预测
拟南芥胚珠共聚焦图像的细胞边界预测
拟南芥侧根光片显微图像的细胞核预测

这些示例不仅展示了模型的实际应用效果，还提供了预训练模型权重，方便用户快速上手和迁移学习。

结语

PyTorch-3DUNet项目为3D体积数据的语义分割任务提供了一个全面而强大的解决方案。无论是在医学影像分析、生物学研究还是其他需要处理3D数据的领域，这个项目都能提供宝贵的工具和资源。我们鼓励感兴趣的研究者和开发者深入探索这个项目，并将其应用到自己的研究或实际问题中。

通过本文的介绍，相信读者已经对PyTorch-3DUNet项目有了全面的了解。无论您是刚接触3D图像分割的新手，还是寻求高效工具的经验丰富的研究者，PyTorch-3DUNet都值得一试。让我们一起期待这个项目在未来带来更多令人兴奋的发展和应用！

🔗 项目GitHub链接

📚 参考文献：

Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., & Ronneberger, O. (2016). 3D U-Net: learning dense volumetric segmentation from sparse annotation. In International conference on medical image computing and computer-assisted intervention (pp. 424-432). Springer, Cham.
Lee, K., Zung, J., Li, P., Jain, V., & Seung, H. S. (2017). Superhuman accuracy on the SNEMI3D connectomics challenge. arXiv preprint arXiv:1706.00120.
Hu, J., Shen, L., & Sun, G. (2018). Squeeze-and-excitation networks. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 7132-7141).