U-Net Brain Tumor Segmentation
This repo show you how to train a U-Net for brain tumor segmentation. By default, you need to download the training set of BRATS 2017 dataset, which have 210 HGG and 75 LGG volumes, and put the data folder along with all scripts.
data
-- Brats17TrainingData
-- train_dev_all
model.py
train.py
...
### About the data Note that according to the license, user have to apply the dataset from BRAST, please do **NOT** contact me for the dataset. Many thanks.
Fig 1: Brain Image
Label 0: background
Label 1: necrotic and non-enhancing tumor
Label 2: edema
Label 4: enhancing tumor
The prepare_data_with_valid.py
split the training set into 2 folds for training and validating. By default, it will use only half of the data for the sake of training speed, if you want to use all data, just change DATA_SIZE = 'half'
to all
.
About the method
Fig 2: Data augmentation
Start training
We train HGG and LGG together, as one network only have one task, set the task
to all
, necrotic
, edema
or enhance
, "all" means learn to segment all tumors.
python train.py --task=all
Note that, if the loss stick on 1 at the beginning, it means the network doesn't converge to near-perfect accuracy, please try restart it.