This project aims at improving U-Net for medical images segmentation. Our model was implemented using Tensorflow and Keras, and the CRF-RNN layer refers to this repo
- U-Net with CRF-RNN layer paper:
- Reference paper:
This repo provides an U-Net with the CRF-RNN layer, and also provides some extract models for comparison, like SegNet, FCN, vanilla U-Net and so on.
modelFns = {'unet':Models.VanillaUnet.VanillaUnet,
'segnet':Models.Segnet.Segnet ,
'vgg_unet':Models.VGGUnet.VGGUnet ,
'vgg_unet2':Models.VGGUnet.VGGUnet2 ,
'fcn8':Models.FCN8.FCN8,
'fcn32':Models.FCN32.FCN32,
'crfunet':Models.CRFunet.CRFunet }
- data hierarchy
Use the Keras data generators to load train and test
Image and label are in structure:
train/
img/
0/
gt/
0/
test/
img/
0/
gt/
0/
- Training parameters
'--batch_size', type=int, default=1, help='input batch size'
'--learning_rate', type=float, default=0.0001, help='learning rate'
'--lr_decay', type=float, default=0.9, help='learning rate decay'
'--epoch', type=int, default=80, help='# of epochs'
'--imSize', type=int, default=320, help='then crop to this size'
'--iter_epoch', type=int, default=0, help='# of iteration as an epoch'
'--num_class', type=int, default=2, help='# of classes'
'--checkpoint_path', type=str, default='', help='where checkpoint saved'
'--data_path', type=str, default='', help='where dataset saved. See loader.py to know how to organize the dataset folder'
'--load_from_checkpoint', type=str, default='', help='where checkpoint saved'
- Train your model
python train.py --data_path ./datasets/ --checkpoint_path ./checkpoints/ - Visualize the train loss, dice score, learning rate, output mask, and first layer convolutional kernels per iteration in tensorboard
tensorboard tensorboard --logdir=./checkpoints - Evaluate your model
python eval.py --data_path ./datasets/ --load_from_checkpoint ./checkpoints/model-xxxx
- Dataset
- Parameters
| param | value |
|---|---|
| batch_size | 5 |
| epoch | 80 |
| iter_epoch | 10 |
| imSize | 320 |
| learning_rate | 0.001 |
| lr_decay | 0.9 |
- Result
| model | IU | DSC | PA |
|---|---|---|---|
| CNN-CRF | 68.73% | 73.22% | 51.77% |
| FCN-8s | 59.61% | 65.73% | 44.26% |
| Segnet | 70.85% | 79.01% | 58.03% |
| Vanilla U-Net | 75.42% | 83.49% | 72.18% |
| U-Net-CRF | 78.00% | 85.77% | 79.05% |
| Our method | 79.89% | 87.31% | 81.27% |
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent