Search code examples
pythontensorflowkeraskeras-layer

Shape error when using CRF for binary segmentation in keras


I am trying to apply a CRF layer to my network for binary segmentation, but it is resulting in an error ValueError: Shapes (?, 1, 1) and (?, 336, 1) are not compatible.

I want to output a tensor with the shape (batch_size, 224, 336, 1). Based on the error, the img_height appears to have been lost once the CRF was introduced.

Below is some sample code describing the model. Without the CRF at the end it works well.

import keras
from keras.layers import UpSampling2D, Conv2D, Activation, MaxPooling2D
from keras_contrib.layers import CRF

img_width, img_height = 336, 224
kernel_size = 7

input=keras.engine.topology.Input(shape=(img_height, img_width, 3))

e=Conv2D(32,(kernel_size,kernel_size),padding='same')(input)
e1=Activation('relu')(e)
e=MaxPooling2D(pool_size=(2, 2))(e1)

e=Conv2D(64,(kernel_size,kernel_size),padding='same')(e)
e2=Activation('relu')(e)
e=MaxPooling2D(pool_size=(2, 2))(e2)

#Decoder layers
d=UpSampling2D()(e)
d=Conv2D(64,(kernel_size,kernel_size),padding='same')(d)
d=Activation('relu')(d)

d=UpSampling2D()(d)
d=Conv2D(32,(kernel_size,kernel_size),padding='same')(d)
d=Activation('relu')(d)

d=Conv2D(1,(1,1),padding='valid')(d)
d=Activation('sigmoid')(d)

out=CRF(1, sparse_target=True)(d) 

autoencoder = Model(inputs=input, outputs=out)

What is the correct way to add the CRF to my segmentation network?


Solution

  • It turns out that keras_contrib.layers.CRF is meant only for sequential data, not spatial data. For using spatial data, the densecrf library works. It appears that it cannot be used during training, only for post-processing. I used a tutorial to implement this, which can be found here: http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/18/image-segmentation-with-tensorflow-using-cnns-and-conditional-random-fields/