用于使用PyTorch深度学习模型进行体数据语义分割的工具包
项目描述
开发正在迁移
该包的开发正在迁移到Rosalind Franklin研究所。现在可以在https://github.com/rosalindfranklininstitute/volume-segmantics找到分支。
Volume Segmantics
使用PyTorch深度学习模型进行体数据语义分割的工具包。
体积语义学提供了一种简单的命令行界面和API,允许研究人员快速训练多种2D PyTorch分割模型(例如U-Net、U-Net++、FPN、DeepLabV3+)在其3D数据集上。这些模型使用预训练的编码器,使得在小数据集上快速训练成为可能。随后,该库允许使用这些训练好的模型来分割更大的3D数据集,自动合并正交平面和旋转中做出的预测,以减少使用2D网络进行3D分割预测可能产生的伪影。
给定一个3D图像体积和相应的密集标签(分割),在沿着x、y和z轴取出的图像切片上训练一个2D模型。该方法针对小训练数据集进行了优化,例如在$128^3$和$512^3$像素之间的单个数据集。为此,所有模型都使用预训练的编码器,并使用图像增强来扩展训练数据集的大小。
这项工作利用了优秀的segmentation-models-pytorch库提供的功能,并结合Albumentations提供的增强功能。此外,所使用的指标和损失函数也利用了Adrian Wolny在pytorch-3dunet存储库中所做的辛勤工作。
要求
需要一台能够运行CUDA启用PyTorch 1.7.1或更高版本的机器。这通常意味着一个相当现代的NVIDIA GPU。具体要求因操作系统而异。例如,在Windows上,您将需要Visual Studio Build Tools以及已安装的CUDA Toolkit,更多详细信息请参阅CUDA文档。
安装
安装该包最简单的方法是首先创建一个新的conda环境或virtualenv,并包含python(理想情况下版本>=3.8)和pip,然后激活该环境并执行pip install volume-segmantics
。如果pip没有安装CUDA启用的PyTorch构建,您可以尝试pip install volume-segmantics --extra-index-url https://download.pytorch.org/whl
,这似乎是Windows上的一个问题。
配置和命令行使用
安装后,当您的环境激活时,您的终端将提供两个新命令:model-train-2d
和model-predict-2d
。
这些命令需要访问存储在YAML文件中的某些设置。这些文件需要位于您运行命令的目录中的名为volseg-settings
的目录中。设置文件可以从此处复制。
可以编辑文件2d_model_train_settings.yaml
以更改训练参数,例如训练轮数、损失函数、评估指标以及模型和编码器架构。可以编辑文件2d_model_predict_settings.yaml
以更改参数,例如预测“质量”,例如“低”质量表示通过沿单个轴((x,y)平面中的图像)提取图像来预测体积分割。对于“中”和“高”质量,分别沿三个轴和12个方向(三个轴,四个旋转)进行预测,然后通过最大概率进行组合。
在3D图像体积及其相应标签上训练2D模型
运行以下命令。输入文件可以是HDF5或多页TIFF格式。
model-train-2d --data path/to/image/data.h5 --labels path/to/corresponding/segmentation/labels.h5
可以在--data
和--labels
标志之后添加多个数据体积和标签体积的路径。将根据/volseg-settings/2d_model_train_settings.yaml
中定义的设置训练一个模型,并将其保存到您的当前工作目录。此外,还会保存一个显示验证集中一些图像的“真实”分割与模型分割的图像。
使用2D模型进行3D体积分割预测
运行以下命令。输入图像文件可以是HDF5或多页TIFF格式。
model-predict-2d path/to/model_file.pytorch path/to/data_for_prediction.h5
输入数据将使用输入模型进行分割,遵循在volseg-settings/2d_model_predict_settings.yaml
中指定的设置。一个包含分割体积的HDF5文件将被保存到您的当前工作目录。
使用示例数据教程
这里有一个教程在此,它提供了如何从收集在人类胎盘组织样本上的同步辐射X射线微CT数据中分割血管的步骤。
目前支持的模式架构和解码器
目前可用并已测试的模式架构如下
- U-Net
- U-Net++
- FPN
- DeepLabV3
- DeepLabV3+
- MA-Net
- LinkNet
- PAN
可用于这些架构的预训练解码器包括
- ResNet-34
- ResNet50
- ResNeXt-50_32x4d
- Efficientnet-b3
- Efficientnet-b4
- Resnest50d*
- Resnest101e*
* 带有星号的解码器与PAN不兼容。
使用API
您可以通过API在自己的程序中使用此包的功能,这在此处有文档说明[链接]。此接口是SuRVoS2使用的,这是一个客户端/服务器GUI应用程序,允许快速注释和分割体积数据。
贡献
我们欢迎社区贡献。请参阅我们的贡献指南以获取更多信息。
引用
如果您使用此包进行研究,请引用以下内容
@article{King2022,
doi = {10.21105/joss.04691},
url = {https://doi.org/10.21105/joss.04691},
year = {2022},
publisher = {The Open Journal},
volume = {7},
number = {78},
pages = {4691},
author = {Oliver N. F. King and Dimitrios Bellos and Mark Basham},
title = {Volume Segmantics: A Python Package for Semantic Segmentation of Volumetric Data Using Pre-trained PyTorch Deep Learning Models},
journal = {Journal of Open Source Software} }
参考文献
Albumentations
Buslaev, A.,Iglovikov, V.I.,Khvedchenya, E.,Parinov, A.,Druzhinin, M. 和 Kalinin, A.A. (2020)。Albumentations:快速灵活的图像增强。信息11。 [链接]
Segmentation Models PyTorch
Yakubovskiy, P. (2020)。Segmentation Models Pytorch。 GitHub
PyTorch-3dUnet
Wolny, A.,Cerrone, L.,Vijayan, A.,Tofanelli, R.,Barro, A.V.,Louveaux, M.,Wenzl, C.,Strauss, S.,Wilson-Sánchez, D.,Lymbouridou, R. 等. (2020)。准确且通用的3D细胞分辨率植物组织分割。ELife 9,e57613。 [链接]