稳健分割的 A 层理论视角
这是我们关于 A-Sheaf-Theoretic-Perspective-for-Robust-Segmentation 工作的存储库
描述
该存储库包含用于训练鲁棒分割模型的代码,方法是在离散潜在空间中强制执行形状等变性,并使用细胞束理论对输出分割图拓扑的组合性进行建模,并强制执行基于组合的损失(更多详细信息,请参阅论文)。该代码库包含我们模型的训练和模型代码。我们有不同类型的模型。我们的模型使用基于对比的损失来强制等方差,如我们论文中所述。我们使用对比基础损失来强制二面体群 (D4) 的等变性。我们还通过将模型中的卷积核限制为常规或不可约的组表示来强制等变。
入门
依赖关系
- 请准备一个python=3.7的环境,然后使用命令“pip install -rrequirements.txt”获取依赖项。
数据集
- 您需要创建 3 个 csv 文件(train.csv、validation.csv、test.csv)。 train.csv 应具有三列('t2image'、'adcimage'、't2label'),其中包含图像的路径和相应的分段。 validation.csv 和 test.csv 应有两列('t2image'、't2label'),其中包含图像的路径和相应的分段。我们支持 nifti 格式。我们在 data/Prostate 中提供了前列腺数据的示例。
- 您可以自由选择在您选择的数据集上进行训练,并根据您的意愿进行预处理。我们为前列腺数据集提供了数据加载器。
- 前列腺:前列腺数据集是从 NCI-ISBI13 Challenge 和 decathalon 数据集获取的。
培训/测试。
- 您可以将训练/测试脚本与 main.py 一起运行。您必须输入训练、验证和测试 csv 文件的路径以及保存结果和图像的输出目录。您需要根据您的数据集调整其他超参数,可以在 main.py 中看到。我们有 4 个模型:“ShapeVQUnet”、“HybridShapeVQUnet”、“HybridSE3VQUnet”、“3DSE3VQUnet”。 “ShapeVQUnet”和“HybridShapeVQUnet”模型使用基于对比的损失将潜在空间限制为 D4 组的等变形状空间。如果您选择“ShapeVQUnet”或“HybridShapeVQUnet”模型,则应选择参数 --contrastive True,否则应选择 --contrastive False。 “ShapeVQUnet”是 3D 模型,而“HybridShapeVQUnet”是 2D/3D 模型。 “HybridSE3VQUnet”和“3DSE3VQUnet”模型将卷积核限制为 SE3 组。如果您选择“HybridSE3VQUnet”和“3DSE3VQUnet”之一,则必须选择是否需要常规(“Regular”)或不可约(“Irreducible”)组表示(--repr)。如果您选择常规(“常规”)组表示形式,则您必须选择组 (--group),例如 --group 4 与 D4 组等方差。如果选择“HybridSE3VQUnet”和“3DSE3VQUnet”模型,您还必须选择组中每个元素的重数 (--multiplicity)。对于所有模型,您还必须选择在多少个纪元之后包含基于细胞束的损失 (--topo_epoch) 以下是前列腺数据的示例
python main.py --modeltype 'HybridShapeVQUnet' --contrastive True --topo_epoch 25 --training_data '.../Sheaves_for_Segmentation/data/Prostate/train.csv' --validation_data '.../Sheaves_for_Segmentation/data/Prostate/validation.csv' --test_data '.../Sheaves_for_Segmentation/data/Prostate/test.csv', --output_directory '.../Sheaves_for_Segmentation/data/Prostate/output/'
作者
贡献者姓名和联系信息
艾因卡兰·桑蒂拉塞卡拉姆 ([email protected])
参考