WMH_Segmentation

White matter hyperintensity segmentation to find lesions in the brain.

WMH Segmentation using data available on https://wmh.isi.uu.nl/

Using T1 and FLAIR images for predictions.

40 subjects split in 8:2 train to test split used. First sliced and normalized across volume and trained on 2D slices.

Labels 0, 1, 2 corresponding to background, WMH and other pathologies, respectively. Label to multi-channel transform used to create 2 channel representing WMH and background(merging class 0 and 2).

Model:

UNet(

dimensions=2,
in_channels=2,    
out_channels=2,    
channels=(16, 32, 64, 128, 256),    
strides=(2, 2, 2, 2),    
num_res_units=2,        
dropout=0.2,    
kernel_size=3,

)

Optimizer:

optimizer = torch.optim.Adam(
model.parameters(), 1e-3, weight_decay=1e-5, amsgrad=True,
)

Sample training Predictions at various epochs (Output=Predictions):

Image Channel 0 - T1

Image Channel 1 - FLAIR

Epoch 42

Epoch 334

Epoch 354

Training and validation loss:

Prediction:

8 subject WMH segmentation DICE score
[0.9567131996154785, 0.9833950400352478, 0.9731700420379639, 0.9664475917816162, 0.9873022437095642, 0.9588155746459961, 0.978124737739563, 0.954979658126831]

Mean DICE score
0.96987 (5 d.p.)

Hausdorff distance score:
To implement

Sample subject prediction visualization:

Tasks to do:

• Implement scheduler to improve loss by reducing learning rate over epoch
• Move code into Project folder and separate into relevante modules and classes
• Train on more data (if applicable and relavant for segmentation task in question)
• Predict on more dataset and compare results.