穩健分割的 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])
參考