Modification of convolutional neural net "UNET" for image segmentation in Keras framework
Python 3.*, Keras 2.1, Tensorflow 1.4
from zf_unet_224_model import ZF_UNET_224, dice_coef_loss, dice_coef
from keras.optimizers import Adam
model = ZF_UNET_224(weights='generator')
optim = Adam()
model.compile(optimizer=optim, loss=dice_coef_loss, metrics=[dice_coef])
model.fit(...)
- "ZF_UNET_224" Model based on UNET code from following paper: https://arxiv.org/abs/1505.04597
- This model used to get 2nd place in DSTL competition: https://www.kaggle.com/c/dstl-satellite-imagery-feature-detection
- For training used DICE coefficient: https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient
- Input shape for model is 224x224 (the same as for other popular CNNs like VGG or ResNet)
- It has 3 input channels (to process standard RGB (BGR) images). You can change it with variable "INPUT_CHANNELS"
- In most cases model ZF_UNET_224 is ok to be used without pretrained weights.
- This code should work fine on both Theano and Tensorflow backends. Code prepared for Keras 2.1, if you need code for Keras 1.2 then use this link:
Download: Weights for Tensorflow backend ~123 MB (Keras 2.1, Dice coef: 0.998)
Weights were obtained with random image generator (generator code available here: train_infinite_generator.py). See example of images from generator below.
Dice coefficient for pretrained weights: ~0.998. See history of learning below: