This repo showcases two deep learning methods(SegNet and UNet) with Pytorch to segment fetal images. The work is heavily based on the book PyTorch Computer Vision Cookbook-Michael Avendi.
data/training/: Stores all the unsplited training data
data/test_set: Stores all data for prediction use
model/weights.pt: Stores the best weight model generated after training
main.py: The main script imports dataset, trainer, loss functions to run the model
dataset.py: Customise a dataset to process the trainig images
model.py: Construct the SegNet and UNet model
train.py: The trainer to run epochs
loss_functions.py: Define the dice loss + BCElogits loss function
predict.py: Script to predict unlabeld images
torch == 1.8.0torchvisiontorchsummarynumpyscipyskimagematplotlibPIL
SegNet Model Results
This UNet implementation is rather a vanilla model, there is no BatchNorm, DropOut utilised. If one follow the original paper strictly, there will be a conflict betweent input and output sizes(572 to 388). To avoid label and prediction mismatch in this implementatino, a resize function has been applied after every up-convolution in expansive path and at final output layer.
Unet Model Results
