This is an implementaion of volumetric segmentation of 3D medical images of heart using a standard Unet(Learning Dense Volumetric Segmentation from Sparse Annotation Özgün Çiçek et al.. ) This code can be used for binary and multiclass semantic segmentation of images.
-
Install CUDA
-
Install PyTorch
-
Install dependencies
pip install -r requirements.txt
-
Download the data in dataset/data folder and train the model with respective model parameters.
python train.py --epochs 500 --batch_size 5 --learning_rate 1e-5
-
Predict test data with saved model in models path.
python predict.py --model best_model.pth --input filename
We obtained excellent segmentation results for EM cell images. The loss function converged well in 200 iterations.