from __future__ import division, print_function
import unittest
from keras.layers import Input
from keras import backend as K
from rvseg.models import convunet
from rvseg.models import unet
class TestModel(unittest.TestCase):
def test_downsampling(self):
inputs = Input(shape=(28, 28, 1))
filters = 16
padding = 'valid'
x, y = convunet.downsampling_block(inputs, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 12, 12, filters))
self.assertTupleEqual(K.int_shape(y), (None, 24, 24, filters))
padding = 'same'
x, y = convunet.downsampling_block(inputs, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 14, 14, filters))
self.assertTupleEqual(K.int_shape(y), (None, 28, 28, filters))
def test_downsampling_error(self):
# downsampling should fail on odd-integer dimension images
inputs = Input(shape=(29, 29, 1))
filters = 16
with self.assertRaises(AssertionError):
convunet.downsampling_block(inputs, filters, padding='valid')
with self.assertRaises(AssertionError):
convunet.downsampling_block(inputs, filters, padding='same')
def test_upsampling(self):
# concatenation without cropping
filters = 16
inputs = Input(shape=(14, 14, 2*filters))
skip = Input(shape=(28, 28, filters))
padding = 'valid'
x = convunet.upsampling_block(inputs, skip, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 24, 24, filters))
# ((4,4), (4,4)) cropping
filters = 15
inputs = Input(shape=(10, 10, 2*filters))
skip = Input(shape=(28, 28, filters))
padding = 'valid'
x = convunet.upsampling_block(inputs, skip, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 16, 16, filters))
# odd-integer input size
filters = 4
inputs = Input(shape=(11, 11, 2*filters))
skip = Input(shape=(28, 28, filters))
padding = 'valid'
x = convunet.upsampling_block(inputs, skip, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 18, 18, filters))
# test odd-integer cropping
filters = 5
inputs = Input(shape=(11, 11, 2*filters))
skip = Input(shape=(27, 27, filters))
padding = 'valid'
x = convunet.upsampling_block(inputs, skip, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 18, 18, filters))
# test same padding
filters = 5
inputs = Input(shape=(11, 11, 2*filters))
skip = Input(shape=(27, 27, filters))
padding = 'same'
x = convunet.upsampling_block(inputs, skip, filters, padding)
self.assertTupleEqual(K.int_shape(x), (None, 22, 22, filters))
def test_upsampling_error(self):
filters = 2
inputs = Input(shape=(11, 11, 2*filters))
padding = 'valid'
with self.assertRaises(AssertionError):
skip = Input(shape=(21, 22, filters))
x = convunet.upsampling_block(inputs, skip, filters, padding)
with self.assertRaises(AssertionError):
skip = Input(shape=(22, 21, filters))
x = convunet.upsampling_block(inputs, skip, filters, padding)
def test_unet(self):
# classic u-net architecture from
# "U-Net: Convolutional Networks for Biomedical Image Segmentation"
# O. Ronneberger, P. Fischer, T. Brox (2015)
height, width, channels = 572, 572, 1
features = 64
depth = 4
classes = 2
temperature = 1.0
padding = 'valid'
m = unet(height, width, channels, classes, features, depth,
temperature, padding)
self.assertEqual(len(m.layers), 56)
# input/output dimensions
self.assertTupleEqual(K.int_shape(m.input), (None, 572, 572, 1))
self.assertTupleEqual(K.int_shape(m.output), (None, 388, 388, 2))
# layers
layer_output_dims = [
(None, 572, 572, 1), # input
(None, 570, 570, 64),
(None, 570, 570, 64),
(None, 568, 568, 64), # skip 1
(None, 568, 568, 64),
(None, 284, 284, 64), # max pool 2x2
(None, 282, 282, 128),
(None, 282, 282, 128),
(None, 280, 280, 128), # skip 2
(None, 280, 280, 128),
(None, 140, 140, 128), # max pool 2x2
(None, 138, 138, 256),
(None, 138, 138, 256),
(None, 136, 136, 256), # skip 3
(None, 136, 136, 256),
(None, 68, 68, 256), # max pool 2x2
(None, 66, 66, 512),
(None, 66, 66, 512),
(None, 64, 64, 512), # skip 4
(None, 64, 64, 512),
(None, 32, 32, 512), # max pool 2x2
(None, 30, 30, 1024),
(None, 30, 30, 1024),
(None, 28, 28, 1024),
(None, 28, 28, 1024),
(None, 56, 56, 512), # up-conv 2x2
(None, 56, 56, 512), # cropping of skip 4
(None, 56, 56, 1024), # concat
(None, 54, 54, 512),
(None, 54, 54, 512),
(None, 52, 52, 512),
(None, 52, 52, 512),
(None, 104, 104, 256), # up-conv 2x2
(None, 104, 104, 256), # cropping of skip 3
(None, 104, 104, 512), # concat
(None, 102, 102, 256),
(None, 102, 102, 256),
(None, 100, 100, 256),
(None, 100, 100, 256),
(None, 200, 200, 128), # up-conv 2x2
(None, 200, 200, 128), # cropping of skip 2
(None, 200, 200, 256), # concat
(None, 198, 198, 128),
(None, 198, 198, 128),
(None, 196, 196, 128),
(None, 196, 196, 128),
(None, 392, 392, 64), # up-conv 2x2
(None, 392, 392, 64), # cropping of skip 1
(None, 392, 392, 128), # concat
(None, 390, 390, 64),
(None, 390, 390, 64),
(None, 388, 388, 64),
(None, 388, 388, 64),
(None, 388, 388, 2), # output segmentation map
(None, 388, 388, 2),
(None, 388, 388, 2),
]
for layer, shape in zip(m.layers, layer_output_dims):
self.assertTupleEqual(layer.output_shape, shape)
def check_layer_dims(self, model):
# if we include only one of batch normalization or dropout,
# then the shape of the network should be the same.
layer_output_dims = [
(None, 10, 10, 1), # input
(None, 10, 10, 4), # conv2D
(None, 10, 10, 4), # batchnorm | reLU
(None, 10, 10, 4), # reLU | dropout
(None, 10, 10, 4), # conv2D
(None, 10, 10, 4), # batchnorm | reLU
(None, 10, 10, 4), # reLU | dropout
(None, 5, 5, 4), # max pool 2x2
(None, 5, 5, 8), # conv2D
(None, 5, 5, 8), # batchnorm | reLU
(None, 5, 5, 8), # reLU | dropout
(None, 5, 5, 8), # conv2D
(None, 5, 5, 8), # batchnorm | reLU
(None, 5, 5, 8), # reLU | dropout
(None, 10, 10, 4), # up-conv 2x2
(None, 10, 10, 8), # concat
(None, 10, 10, 4), # conv2D
(None, 10, 10, 4), # batchnorm | reLU
(None, 10, 10, 4), # reLU | dropout
(None, 10, 10, 4), # conv2D
(None, 10, 10, 4), # batchnorm | reLU
(None, 10, 10, 4), # reLU | dropout
(None, 10, 10, 2), # output segmentation map
(None, 10, 10, 2), # (temperature)
(None, 10, 10, 2), # softmax
]
for layer, shape in zip(model.layers, layer_output_dims):
self.assertTupleEqual(layer.output_shape, shape)
def test_batchnorm(self):
# only batch norm, no dropout
height, width, channels = 10, 10, 1
features = 4
depth = 1
classes = 2
temperature = 1.0
padding = 'same'
batchnorm = True
dropout = False
m = unet(height, width, channels, classes, features, depth,
temperature, padding, batchnorm, dropout)
self.assertEqual(len(m.layers), 25)
# input/output dimensions
self.assertTupleEqual(K.int_shape(m.input), (None, 10, 10, 1))
self.assertTupleEqual(K.int_shape(m.output), (None, 10, 10, 2))
self.check_layer_dims(m)
def test_dropout(self):
# only dropout, no batch norm
height, width, channels = 10, 10, 1
features = 4
depth = 1
classes = 2
temperature = 1.0
padding = 'same'
batchnorm = False
dropout = True
m = unet(height, width, channels, classes, features, depth,
temperature, padding, batchnorm, dropout)
self.assertEqual(len(m.layers), 25)
# input/output dimensions
self.assertTupleEqual(K.int_shape(m.input), (None, 10, 10, 1))
self.assertTupleEqual(K.int_shape(m.output), (None, 10, 10, 2))
self.check_layer_dims(m)
