This framework designed for medical image segmentation. It uses the ushaped deep CNN architecture (VNet or UNet) extended by attention gates. For more information about the method, please read the following paper:
https://doi.org/10.3389/frobt.2020.00106
Please also cite this paper if you are using this code for your research!
The attention mechanism helps the network to focus on the desired area of CT abdomen scan and therefore improves the accuracy. The attention maps from two network levels (image resolutions) for three different abdomen CT images datasets are in Fig 1 below.
This repository is still work in progress. Things may break. If that is the case, please let us know. The code was wastly inspired by the nnU-Net framework authored by Fabian Isensee, so in case of troubles you may want to study this repository and its describtion, too: https://github.com/MIC-DKFZ/nnUNet/tree/master/nnunet.
The instalation is only tested on Linux (Ubuntu). It may work on other operating systems as well but we do not guarantee it will.
Installation instructions
- Install PyTorch (https://pytorch.org/get-started/locally/)
- Clone this repository
git clone https://github.com/tureckova/Abdomen-CT-Image-Segmentation
- Go into the repository (
cd Abdomen-CT-Image-Segmentation
on linux) - Install with
pip install -r requirements.txt
followed bypip install -e .
All the commands in this section assume that you are in a terminal and your working directory is the repository folder
(the one that has all the subfolders like dataset_conversion
, evaluation
, ...)
Framwork needs to know where you will store raw data, want it to store preprocessed data and trained models. Have a
look at the file paths.py
and adapt it to your system.
The preprocessing pipeline was adapted from nnU-Net. Please refer to the readme.md in the
dataset_conversion
subfolder for detailed information. Examples are also provided there. You will need to
convert your dataset into this format before you can continue.
Place your dataset either in the raw_dataset_dir
or splitted_4d_output_dir
, as specified in paths.py
(depending on how you prepared it, again
see the readme in dataset_conversion
). Give
it a name like: TaskXX_MY_DATASET
(where XX is some number) to be consistent with the naming scheme of the Medical
Segmentation Decathlon.
Framework can now analyze your dataset and determine how to train its models. To run experiment planning and preprocessing for your dataset, execute the following command:
python3 experiment_planning/plan_and_preprocess_task.py -t Task07_Pancreas -p 8
here TaskXX_MY_DATASET
specifies the task (your dataset) and -p
determines how many processes will be used for
datatset analysis and preprocessing. Generally you want this number to be as high as you have CPU cores, unless you
run into memory problems (beware of datasets such as LiTS!)
Running this command will to several things:
- If you stored your data as 4D nifti the data will be split into a sequence of 3d niftis. Back when I started SimpleITK did not support 4D niftis. This was simply done out of necessity.
- Images are cropped to the nonzero region. This does nothing for most datasets. Most brain datasets however are brain
extracted, meaning that the brain is surrounded by zeros. There is no need to push all these zeros through the GPUs so
we simply save a little bit of time doing this. Cropped data is stored in the
cropped_output_dir
(paths.py
). - The data is analyzed and information about spacing, intensity distributions and shapes are determined
- nnU-Net configures the architectures based on that information. All U-Nets are configured to optimally use 12GB Nvidia TitanX GPUs.
- The preprocessing is run and it saves the preprocessed data and plans files in
preprocessing_output_dir
. You could accomodate the plans files (for example to fit in smaller GPU), the skript rewrite_pickle_plan.py may help you.
I strongly recommend you set preprocessing_output_dir
on a SSD. HDDs are typically too slow for data loading. Plan files as we use it is the paper here in the folder plans_files.
In our paper we compared two variants of 3D CNN - VNet or UNet and two resolution variants - full-resolution and low-resolution. The default setting is to train each of these models in a five-fold cross-validation.
Trained models are stored in network_training_output_dir
(specified in paths.py
).
For FOLD
in [0, 4], run:
python run/run_training.py [3d_fullres/3d_lowres] nnUNetTrainer TaskXX_MY_DATASET FOLD --ndet --vnet=[0/1]
you need to choose only one of variants in [] and remove the brackets.
You can continue the already started training of the model by adding --continue_training to the command. The model will recover from the newest checkpoint.
The model is evaluated automaticaly in the end of the training, but the model could be only evaluated by the same command as training only with added -val
in the end of it.
You can use trained models to predict test data. In order to be able to do so the test data must be provided in the same format as the training data.
To run inference use the following script:
python inference/predict_simple.py -i INPUT_FOLDER -o OUTPUT_FOLDER -t TaskXX_MY_DATASET -tr nnUNetTrainer -m [3d_fullres/3d_lowres] --vnet=[0/1]
If you wish to ensemble different inference cases, run all inference commands with the -z
argument. This will tell the framework to save the softmax probabilities as well. They are needed for ensembling.
You can then ensemble the predictions of two output folders (there must be saved the softmax probabilities, see above) with the following command:
python inference/ensemble_predictions.py -f FOLDER1 FODLER2 ... -o OUTPUT_FOLDER
This will ensemble the predictions located in FODLER1, FOLDER2, ...
and write them into OUTPUT_FOLDER
You can quicly evaluate the esemble files with the folowing command:
python3 inference/validate_nifti_folder.py -f PREDICTED_FOLDER -gtf FOLDER_GROUND_TURE
This will create summary.json file in location PREDICTED_FOLDER with evaluation scores for each file as well as mean evaluation scores.
The cross-validation in nnU-Net splits on a per-case basis. This may sometimes not be desired, for example because
several training cases may be the same patient (different time steps or annotators). If this is the case, then you need to
manually create a split file. To do this, first let nnU-Net create the default split file. Run one of the network
trainings (any of them works fine for this) and abort after the first epoch. nnU-Net will have created a split file automatically:
preprocessing_output_dir/TaskXX_MY_DATASET/splits_final.pkl
. This file contains a list (length 5, one entry per fold).
Each entry in the list is a dictionary with keys 'train' and 'val' pointing to the patientIDs assigned to the sets.
To use your own splits in nnU-Net, you need to edit these entries to what you want them to be and save it back to the
splits_final.pkl file. Use load_pickle and save_pickle from batchgenerators.utilities.file_and_folder_operations for convenience.
You can share trained models by simply sending the corresponding output folder from network_training_output_dir
to
whoever you want share them with. The recipient can then use nnU-Net for inference with this model.