This repository aims to practice pytorch and implement U-net architecture by Ronneberger et al. [1]. The used dataset is nerve ultrasound images and we aim to delineate nerve structures from them.
usage: main.py [-h] [--workers WORKERS] [--batchSize BATCHSIZE]
[--niter NITER] [--start_epoch START_EPOCH] [--lr LR] [--cuda]
[--resume PATH] [--useBN] [--output_name OUTPUT_NAME]
dataroot
positional arguments:
dataroot path to dataset of kaggle ultrasound nerve
segmentation
optional arguments:
-h, --help show this help message and exit
--workers WORKERS number of data loading workers
--batchSize BATCHSIZE
input batch size
--niter NITER number of epochs to train for
--start_epoch START_EPOCH
number of epoch to start
--lr LR learning rate, default=0.0002
--cuda enables cuda
--resume PATH path to latest checkpoint (default: none)
--useBN enalbes batch normalization
--output_name OUTPUT_NAME
output checkpoint filename
To run our pre-trained model, commit:
python3 main.py data --resume checkpoint_BN.tar --niter 0 --useBN
Image
Ground truth
Predicted segmentation