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).
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 = torch.optim.Adam( model.parameters(), 1e-3, weight_decay=1e-5, amsgrad=True, )
Image Channel 0 - T1
Image Channel 1 - FLAIR
Epoch 42
Epoch 334
Epoch 354
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:
- 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.