Repository to our 3d neuro-anatomic segmentation method presented at the MIDL 2022 conference
Roy S*, Kügler D*, Reuter M (*equal contribution). Are 2.5D approaches superior to 3D deep networks in whole brain segmentation? In: Medical Imaging with Deep Learning. Proceedings of Machine Learning Research (2022), PDF, Bibtex
The params dictionary for each model is provided in models/config.py
-
Run
train_model.py
to train a model.CUDA_VISIBLE_DEVICES=0,1 python train_model.py
. 2 GPUs are necessary if you have a separateloss_device
. -
Even though the micro/block architecture file (eg. VNet.py) is copied into the checkpoints folder, it is for reference. During Evaluation the micro/block architectures are taken from ./model/inference. This is to ensure compatibility with older trained models. Fix this if you are changing the micro/block archs.
For evaluation, users should use eval_w_mgz.py
. See documentation in the help.
Name in Paper | Class Name | Source File | Weights File |
---|---|---|---|
Ours | QuadNet | model/QuadNet.py |
ensembled_model_<index> |
Ours (w/o SL) | RCVNet | model/inference/RCVNet.py |
unensembled_model |
While we have included code to automatically download model files, these model files can also be manually downloaded.
These architectures did not provide sufficient performance improvements over the chosen architectures (see above).
Model Description (Backbone VNet with ...) | Source File |
---|---|
GroupNorm and 3^3 Kernels instead of 5^3 | model/RCVNet.py |
Symmetric filter factorization with two 3^3 kernels instead of 5^3 | model/VNetSym.py |
Asymmetric filter factorization with two sets of (3x1x1), (1x3x1), (1x1x3) filters instead of 5^3 | model/VNetAsym.py |
Dense residual block as in DenseNet, with additive residuals | model/VNetDenseAdd.py |
Attention Block as in Attention UNet | model/inference/VNet.py |
Squeeze and Excite block as in SqueezeNet | model/VNetSE.py |
Concurrent 3 axis 2.5D and 3D paths with 75% of feature maps in 3D path | model/VNet_2D_3D.py |
only 75% of feature maps included in each block | model/MultiResVNet.py |
Checkpoint directory:
- Copy the model folder to evaluate into your present working 3dseg's checkpoint folder. Also, change the absolute path to the checkpoints folder to the one inside your present working 3dseg checkpoint folder and create a
__init__.py
(if it does not already exist). This is because model classes will be directly read from here - Go into the right file containing the model macro architecture (as specified by
file_string
) named RCVNet.py / CropNet3d.py / LResFCNet.py inside the copied folder and change the import on line 5 from model.module to .module (basically forcing it to use themodule.py
file inside this folder as the layer/micro architecture definition).module.py
is the common name I gave to VNet.py, VNet_2D_3D.py etc when I copied them over during training, e.g.:- OLD:
from model.module import EncoderBlock, BottleNeck, DecoderBlock
- NEW:
from .module import EncoderBlock, BottleNeck, DecoderBlock
- OLD:
- Select the correct file and class name in the eval*.py file and run.
- There is no named retrieval of classes for quadnet eval. Please make sure that a standard vnet is retrieved from the
./models
folder. Use the models in the./models/inference
folder if necessary - I created that folder to keep copies incase one edited the py files inside the models folder. - Please find details of the model folders in the 3dseg-saved-models.xlsx file