U-net for image segmentation
TensorLayer and TensorFlow implementation of U-Net image segmenation.
The u-net is convolutional network architecture for fast and precise segmentation of images. Up to now it has outperformed the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. It has won the Grand Challenge for Computer-Automated Detection of Caries in Bitewing Radiography at ISBI 2015, and it has won the Cell Tracking Challenge at ISBI 2015 on the two most challenging transmitted light microscopy categories (Phase contrast and DIC microscopy) by a large margin (See also our annoucement).
More details can be found in https://github.com/zsdonghao?tab=repositories
With Simplified Convolutional Layer APIs (TensorLayer 1.2.5)
def u_net_2d_64_1024_deconv(x, n_out=2):from tensorlayer.layers import *
nx = int(x._shape[1])
ny = int(x._shape[2])
nz = int(x._shape[3])
print(" * Input: size of image: %d %d %d" % (nx, ny, nz))
w_init = tf.truncated_normal_initializer(stddev=0.01)
b_init = tf.constant_initializer(value=0.0)
inputs = InputLayer(x, name='inputs')
conv1 = Conv2d(inputs, 64, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv1_1')
conv1 = Conv2d(conv1, 64, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv1_2')
pool1 = MaxPool2d(conv1, (2, 2), padding='SAME', name='pool1')
conv2 = Conv2d(pool1, 128, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv2_1')
conv2 = Conv2d(conv2, 128, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv2_2')
pool2 = MaxPool2d(conv2, (2, 2), padding='SAME', name='pool2')
conv3 = Conv2d(pool2, 256, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv3_1')
conv3 = Conv2d(conv3, 256, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv3_2')
pool3 = MaxPool2d(conv3, (2, 2), padding='SAME', name='pool3')
conv4 = Conv2d(pool3, 512, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv4_1')
conv4 = Conv2d(conv4, 512, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv4_2')
pool4 = MaxPool2d(conv4, (2, 2), padding='SAME', name='pool4')
conv5 = Conv2d(pool4, 1024, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv5_1')
conv5 = Conv2d(conv5, 1024, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='conv5_2')
print(" * After conv: %s" % conv5.outputs)
up4 = DeConv2d(conv5, 512, (3, 3), out_size = (nx/8, ny/8), strides = (2, 2),
padding = 'SAME', act=None, W_init=w_init, b_init=b_init, name='deconv4')
up4 = ConcatLayer([up4, conv4], concat_dim=3, name='concat4')
conv4 = Conv2d(up4, 512, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv4_1')
conv4 = Conv2d(conv4, 512, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv4_2')
up3 = DeConv2d(conv4, 256, (3, 3), out_size = (nx/4, ny/4), strides = (2, 2),
padding = 'SAME', act=None, W_init=w_init, b_init=b_init, name='deconv3')
up3 = ConcatLayer([up3, conv3], concat_dim=3, name='concat3')
conv3 = Conv2d(up3, 256, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv3_1')
conv3 = Conv2d(conv3, 256, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv3_2')
up2 = DeConv2d(conv3, 128, (3, 3), out_size = (nx/2, ny/2), strides = (2, 2),
padding = 'SAME', act=None, W_init=w_init, b_init=b_init, name='deconv2')
up2 = ConcatLayer([up2, conv2] ,concat_dim=3, name='concat2')
conv2 = Conv2d(up2, 128, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv2_1')
conv2 = Conv2d(conv2, 128, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv2_2')
up1 = DeConv2d(conv2, 64, (3, 3), out_size = (nx/1, ny/1), strides = (2, 2),
padding = 'SAME', act=None, W_init=w_init, b_init=b_init, name='deconv1')
up1 = ConcatLayer([up1, conv1] ,concat_dim=3, name='concat1')
conv1 = Conv2d(up1, 64, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv1_1')
conv1 = Conv2d(conv1, 64, (3, 3), act=tf.nn.relu, padding='SAME', W_init=w_init, b_init=b_init, name='uconv1_2')
conv1 = Conv2d(conv1, n_out, (1, 1), act=None, name='uconv1')
print(" * Output: %s" % conv1.outputs)
outputs = tl.act.pixel_wise_softmax(conv1.outputs)
return conv1, outputs
With Professional Convolutional Layer APIs.
def u_net_2d_64_1024_deconv_pro(x, n_out=2):
""" 2-D U-Net for Image Segmentation.
Parameters
-----------
x : tensor or placeholder of input with shape of [batch_size, row, col, channel]
batch_size : int, batch size
n_out : int, number of output channel, default is 2 for foreground and background (binary segmentation)
Returns
--------
network : TensorLayer layer class with identity output
outputs : tensor, the output with pixel-wise softmax
Notes
-----
- Recommend to use Adam with learning rate of 1e-5
"""
batch_size = int(x._shape[0])
nx = int(x._shape[1])
ny = int(x._shape[2])
nz = int(x._shape[3])
print(" * Input: size of image: %d %d %d" % (nx, ny, nz))
## define initializer
w_init = tf.truncated_normal_initializer(stddev=0.01)
b_init = tf.constant_initializer(value=0.0)
## u-net model
# convolution
# with tf.device('\gpu:0'):
net_in = tl.layers.InputLayer(x, name='input')
conv1 = tl.layers.Conv2dLayer(net_in, act=tf.nn.relu,
shape=[3,3,nz,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1')
conv2 = tl.layers.Conv2dLayer(conv1, act=tf.nn.relu,
shape=[3,3,64,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2')
pool1 = tl.layers.PoolLayer(conv2, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool1')
conv3 = tl.layers.Conv2dLayer(pool1, act=tf.nn.relu,
shape=[3,3,64,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3')
conv4 = tl.layers.Conv2dLayer(conv3, act=tf.nn.relu,
shape=[3,3,128,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4')
pool2 = tl.layers.PoolLayer(conv4, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool2')
conv5 = tl.layers.Conv2dLayer(pool2, act=tf.nn.relu,
shape=[3,3,128,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv5')
conv6 = tl.layers.Conv2dLayer(conv5, act=tf.nn.relu,
shape=[3,3,256,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv6')
pool3 = tl.layers.PoolLayer(conv6, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool3')
conv7 = tl.layers.Conv2dLayer(pool3, act=tf.nn.relu,
shape=[3,3,256,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv7')
conv8 = tl.layers.Conv2dLayer(conv7, act=tf.nn.relu,
shape=[3,3,512,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv8')
# print(conv8.outputs) # (10, 30, 30, 512)
pool4 = tl.layers.PoolLayer(conv8, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool,name='pool4')
conv9 = tl.layers.Conv2dLayer(pool4, act=tf.nn.relu,
shape=[3,3,512,1024], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv9')
conv10 = tl.layers.Conv2dLayer(conv9, act=tf.nn.relu,
shape=[3,3,1024,1024], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv10')
print(" * After conv: %s" % conv10.outputs) # (batch_size, 32, 32, 1024)
# deconvoluation
deconv1 = tl.layers.DeConv2dLayer(conv10, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,512,1024], strides=[1,2,2,1], output_shape=[batch_size,nx/8,ny/8,512],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon1_1')
# print(deconv1.outputs) #(10, 30, 30, 512)
deconv1_2 = tl.layers.ConcatLayer([conv8, deconv1], concat_dim=3, name='concat1_2')
deconv1_3 = tl.layers.Conv2dLayer(deconv1_2, act=tf.nn.relu,
shape=[3,3,1024,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1_3')
deconv1_4 = tl.layers.Conv2dLayer(deconv1_3, act=tf.nn.relu,
shape=[3,3,512,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1_4')
deconv2 = tl.layers.DeConv2dLayer(deconv1_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,256,512], strides=[1,2,2,1], output_shape=[batch_size,nx/4,ny/4,256],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon2_1')
deconv2_2 = tl.layers.ConcatLayer([conv6, deconv2], concat_dim=3, name='concat2_2')
deconv2_3 = tl.layers.Conv2dLayer(deconv2_2, act=tf.nn.relu,
shape=[3,3,512,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2_3')
deconv2_4 = tl.layers.Conv2dLayer(deconv2_3, act=tf.nn.relu,
shape=[3,3,256,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2_4')
deconv3 = tl.layers.DeConv2dLayer(deconv2_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,128,256], strides=[1,2,2,1], output_shape=[batch_size,nx/2,ny/2,128],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon3_1')
deconv3_2 = tl.layers.ConcatLayer([conv4, deconv3], concat_dim=3, name='concat3_2')
deconv3_3 = tl.layers.Conv2dLayer(deconv3_2, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,256,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3_3')
deconv3_4 = tl.layers.Conv2dLayer(deconv3_3, act=tf.nn.relu,
shape=[3,3,128,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3_4')
deconv4 = tl.layers.DeConv2dLayer(deconv3_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,64,128], strides=[1,2,2,1], output_shape=[batch_size,nx,ny,64],
padding='SAME', W_init=w_init, b_init=b_init, name='devconv4_1')
deconv4_2 = tl.layers.ConcatLayer([conv2, deconv4], concat_dim=3, name='concat4_2')
deconv4_3 = tl.layers.Conv2dLayer(deconv4_2, act=tf.nn.relu,
shape=[3,3,128,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_3')
deconv4_4 = tl.layers.Conv2dLayer(deconv4_3, act=tf.nn.relu,
shape=[3,3,64,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_4')
network = tl.layers.Conv2dLayer(deconv4_4,
act=tf.identity,
shape=[1,1,64,n_out], # [0]:foreground prob; [1]:background prob
strides=[1,1,1,1],
padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_5')
# compute the softmax output
print(" * Output: %s" % network.outputs)
outputs = tl.act.pixel_wise_softmax(network.outputs)
return network, outputs
#
With Simplified Convolutional Layer APIs
def u_net_2d_64_1024_deconv_pro(x, n_out=2):
""" 2-D U-Net for Image Segmentation.
Parameters
-----------
x : tensor or placeholder of input with shape of [batch_size, row, col, channel]
batch_size : int, batch size
n_out : int, number of output channel, default is 2 for foreground and background (binary segmentation)
Returns
--------
network : TensorLayer layer class with identity output
outputs : tensor, the output with pixel-wise softmax
Notes
-----
- Recommend to use Adam with learning rate of 1e-5
"""
batch_size = int(x._shape[0])
nx = int(x._shape[1])
ny = int(x._shape[2])
nz = int(x._shape[3])
print(" * Input: size of image: %d %d %d" % (nx, ny, nz))
## define initializer
w_init = tf.truncated_normal_initializer(stddev=0.01)
b_init = tf.constant_initializer(value=0.0)
## u-net model
# convolution
# with tf.device('\gpu:0'):
net_in = tl.layers.InputLayer(x, name='input')
conv1 = tl.layers.Conv2dLayer(net_in, act=tf.nn.relu,
shape=[3,3,nz,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1')
conv2 = tl.layers.Conv2dLayer(conv1, act=tf.nn.relu,
shape=[3,3,64,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2')
pool1 = tl.layers.PoolLayer(conv2, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool1')
conv3 = tl.layers.Conv2dLayer(pool1, act=tf.nn.relu,
shape=[3,3,64,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3')
conv4 = tl.layers.Conv2dLayer(conv3, act=tf.nn.relu,
shape=[3,3,128,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4')
pool2 = tl.layers.PoolLayer(conv4, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool2')
conv5 = tl.layers.Conv2dLayer(pool2, act=tf.nn.relu,
shape=[3,3,128,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv5')
conv6 = tl.layers.Conv2dLayer(conv5, act=tf.nn.relu,
shape=[3,3,256,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv6')
pool3 = tl.layers.PoolLayer(conv6, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool, name='pool3')
conv7 = tl.layers.Conv2dLayer(pool3, act=tf.nn.relu,
shape=[3,3,256,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv7')
conv8 = tl.layers.Conv2dLayer(conv7, act=tf.nn.relu,
shape=[3,3,512,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv8')
# print(conv8.outputs) # (10, 30, 30, 512)
pool4 = tl.layers.PoolLayer(conv8, ksize=[1,2,2,1],
strides=[1,2,2,1], padding='SAME',
pool=tf.nn.max_pool,name='pool4')
conv9 = tl.layers.Conv2dLayer(pool4, act=tf.nn.relu,
shape=[3,3,512,1024], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv9')
conv10 = tl.layers.Conv2dLayer(conv9, act=tf.nn.relu,
shape=[3,3,1024,1024], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv10')
print(" * After conv: %s" % conv10.outputs) # (batch_size, 32, 32, 1024)
# deconvoluation
deconv1 = tl.layers.DeConv2dLayer(conv10, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,512,1024], strides=[1,2,2,1], output_shape=[batch_size,nx/8,ny/8,512],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon1_1')
# print(deconv1.outputs) #(10, 30, 30, 512)
deconv1_2 = tl.layers.ConcatLayer([conv8, deconv1], concat_dim=3, name='concat1_2')
deconv1_3 = tl.layers.Conv2dLayer(deconv1_2, act=tf.nn.relu,
shape=[3,3,1024,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1_3')
deconv1_4 = tl.layers.Conv2dLayer(deconv1_3, act=tf.nn.relu,
shape=[3,3,512,512], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv1_4')
deconv2 = tl.layers.DeConv2dLayer(deconv1_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,256,512], strides=[1,2,2,1], output_shape=[batch_size,nx/4,ny/4,256],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon2_1')
deconv2_2 = tl.layers.ConcatLayer([conv6, deconv2], concat_dim=3, name='concat2_2')
deconv2_3 = tl.layers.Conv2dLayer(deconv2_2, act=tf.nn.relu,
shape=[3,3,512,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2_3')
deconv2_4 = tl.layers.Conv2dLayer(deconv2_3, act=tf.nn.relu,
shape=[3,3,256,256], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv2_4')
deconv3 = tl.layers.DeConv2dLayer(deconv2_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,128,256], strides=[1,2,2,1], output_shape=[batch_size,nx/2,ny/2,128],
padding='SAME', W_init=w_init, b_init=b_init, name='devcon3_1')
deconv3_2 = tl.layers.ConcatLayer([conv4, deconv3], concat_dim=3, name='concat3_2')
deconv3_3 = tl.layers.Conv2dLayer(deconv3_2, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,256,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3_3')
deconv3_4 = tl.layers.Conv2dLayer(deconv3_3, act=tf.nn.relu,
shape=[3,3,128,128], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv3_4')
deconv4 = tl.layers.DeConv2dLayer(deconv3_4, act=tf.identity, #act=tf.nn.relu,
shape=[3,3,64,128], strides=[1,2,2,1], output_shape=[batch_size,nx,ny,64],
padding='SAME', W_init=w_init, b_init=b_init, name='devconv4_1')
deconv4_2 = tl.layers.ConcatLayer([conv2, deconv4], concat_dim=3, name='concat4_2')
deconv4_3 = tl.layers.Conv2dLayer(deconv4_2, act=tf.nn.relu,
shape=[3,3,128,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_3')
deconv4_4 = tl.layers.Conv2dLayer(deconv4_3, act=tf.nn.relu,
shape=[3,3,64,64], strides=[1,1,1,1], padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_4')
network = tl.layers.Conv2dLayer(deconv4_4,
act=tf.identity,
shape=[1,1,64,n_out], # [0]:foreground prob; [1]:background prob
strides=[1,1,1,1],
padding='SAME',
W_init=w_init, b_init=b_init, name='conv4_5')
# compute the softmax output
print(" * Output: %s" % network.outputs)
outputs = tl.act.pixel_wise_softmax(network.outputs)
return network, outputs