NOTE : This repository is not used anymore. Please check the latest code in here.
This repository provides everything necessary to train and evaluate a brain tumor segmentation model.
The baseline network is Modified-UNet.
Requirements:
- Python 3 (code has been tested on Python 3.5.6)
- PyTorch (code tested with 1.1.0)
- CUDA and cuDNN (tested with Cuda 10.0)
- Python pacakges : tqdm, opencv-python (4.1), SimpleITK (1.2.0), scipy (1.2.1), imgaug (0.4.0)
Structure:
data/
: save directory of datasetsdatasets/
: data loading codenetwork/
: network architecture definitionsoptions/
: argument parser optionsutils/
: image processing code, miscellaneous helper functions, and training/evaluation core codetrain.py/
: code for model trainingevaluate.py/
: code for model evaluationinference.py/
: code for model inference to generate predicted mask nifti files.preprocess_mask.py/
: code for pre-processing masks (refinement of ce mask and peri-tumoral mask generation)
Out private dataset which has four types of MRI images (FLAIR, T1GD, T1, T2) and three types of mask (necro, ce, T2) divided into train (N=139) and test (N=16) dataset.
Place the dataset in data/
directory and the dataset architecture must be as below.
data
└─── train
│ └─── patientDir001
│ │ │ FLAIR_stripped.nii.gz
│ │ │ T1GD_stripped.nii.gz
│ │ │ T1_stripped.nii.gz
│ │ │ T2_stripped.nii.gz
│ │ │ necro_mask.nii.gz
│ │ │ ce_mask.nii.gz
│ │ │ t2_mask.nii.gz
│ │ │
│ └─── patientDir002
│ └─── patientDir003
│ └─── ...
│
└─── valid
└─── patientDir00a
└─── patientDir00b
└─── ...
If the mask doesn't have to be pre-processed (peri-tumoral mask is provided), the patient directory can consist of as below.
patientDir
│ FLAIR_stripped.nii.gz
│ T1GD_stripped.nii.gz
│ T1_stripped.nii.gz
│ T2_stripped.nii.gz
│ necro_mask.nii.gz
│ ce_mask.nii.gz
│ peri_mask.nii.gz
-
Before training, call
python preprocess_mask.py
to pre-process masks.
This python script generatesce_refined_mask.nii.gz
andperi_mask.nii.gz
in each patient directory. -
To train a 3D network, call:
python train.py --batch_size 1 --in_dim 3 --in_depth 128 --in_res 140
-
Before 2D training, call
python parsing_2D.py
to parse 2D datasets, which generates2D_slice/
directory in each patient directory.
To train a 2D network, call:python train.py --batch_size 1 --in_dim 2 --in_res 140
-
To evaluate a network after training, call:
python evaluate.py --in_dim 2 --resume trained_weights.pth
-
To inference a network, call:
python inference.py --in_dim 2 --resume trained_weights.pth
Note that the data composition must be as below and the generated masks are saved in each patient directory.\
data
└─── test
│ └─── patientDir001
│ │ │ FLAIR_stripped.nii.gz
│ │ │ T1GD_stripped.nii.gz
│ │ │ T1_stripped.nii.gz
│ │ │ T2_stripped.nii.gz
│ │ │
│ └─── patientDir002
│ └─── patientDir003
│ └─── ...
Model | Dice_Necro | Dice_CE | Dice_Peri | Dice_Total |
---|---|---|---|---|
M-Unet 2D | 0.9436 | 0.9367 | 0.8658 | 0.9154 |
M-Unet 3D |
- Modified UNet 2D : Google Drive Link
- Modified UNet 3D : Not released yet.
- In 2D training, the intensity scale of input images is changed by scipy resize function in range [0, 255].
The resize function should be replaced to keep the pixel value distribution. (Or the mean and std have to be re-calculated.)
- Release test code.
- Data Augmentation.
- Inference code which generates mask nifti files.
- Release pre-trained models.
- Additional evaluation metrics (Sensitivity, Specificity, Hausdorff95).