| 13 | |
| 14 | # start-snippet-downsampling |
| 15 | def build_UNet(n_input_channels=1, BATCH_SIZE=None, num_output_classes=2, pad='same', nonlinearity=lasagne.nonlinearities.elu, input_dim=(None, None), base_n_filters=64, do_dropout=False): |
| 16 | net = OrderedDict() |
| 17 | net['input'] = InputLayer((BATCH_SIZE, n_input_channels, input_dim[0], input_dim[1])) |
| 18 | |
| 19 | net['contr_1_1'] = batch_norm(ConvLayer(net['input'], base_n_filters, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 20 | net['contr_1_2'] = batch_norm(ConvLayer(net['contr_1_1'], base_n_filters, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 21 | net['pool1'] = Pool2DLayer(net['contr_1_2'], 2) |
| 22 | |
| 23 | net['contr_2_1'] = batch_norm(ConvLayer(net['pool1'], base_n_filters*2, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 24 | net['contr_2_2'] = batch_norm(ConvLayer(net['contr_2_1'], base_n_filters*2, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 25 | net['pool2'] = Pool2DLayer(net['contr_2_2'], 2) |
| 26 | |
| 27 | net['contr_3_1'] = batch_norm(ConvLayer(net['pool2'], base_n_filters*4, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 28 | net['contr_3_2'] = batch_norm(ConvLayer(net['contr_3_1'], base_n_filters*4, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 29 | net['pool3'] = Pool2DLayer(net['contr_3_2'], 2) |
| 30 | |
| 31 | net['contr_4_1'] = batch_norm(ConvLayer(net['pool3'], base_n_filters*8, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 32 | net['contr_4_2'] = batch_norm(ConvLayer(net['contr_4_1'], base_n_filters*8, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 33 | l = net['pool4'] = Pool2DLayer(net['contr_4_2'], 2) |
| 34 | # end-snippet-downsampling |
| 35 | |
| 36 | # start-snippet-bottleneck |
| 37 | # the paper does not really describe where and how dropout is added. Feel free to try more options |
| 38 | if do_dropout: |
| 39 | l = DropoutLayer(l, p=0.4) |
| 40 | |
| 41 | net['encode_1'] = batch_norm(ConvLayer(l, base_n_filters*16, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 42 | net['encode_2'] = batch_norm(ConvLayer(net['encode_1'], base_n_filters*16, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 43 | # end-snippet-bottleneck |
| 44 | |
| 45 | # start-snippet-upsampling |
| 46 | net['upscale1'] = batch_norm(Deconv2DLayer(net['encode_2'], base_n_filters*16, 2, 2, crop="valid", nonlinearity=nonlinearity, W=HeNormal(gain="relu"))) |
| 47 | net['concat1'] = ConcatLayer([net['upscale1'], net['contr_4_2']], cropping=(None, None, "center", "center")) |
| 48 | net['expand_1_1'] = batch_norm(ConvLayer(net['concat1'], base_n_filters*8, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 49 | net['expand_1_2'] = batch_norm(ConvLayer(net['expand_1_1'], base_n_filters*8, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 50 | |
| 51 | net['upscale2'] = batch_norm(Deconv2DLayer(net['expand_1_2'], base_n_filters*8, 2, 2, crop="valid", nonlinearity=nonlinearity, W=HeNormal(gain="relu"))) |
| 52 | net['concat2'] = ConcatLayer([net['upscale2'], net['contr_3_2']], cropping=(None, None, "center", "center")) |
| 53 | net['expand_2_1'] = batch_norm(ConvLayer(net['concat2'], base_n_filters*4, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 54 | net['expand_2_2'] = batch_norm(ConvLayer(net['expand_2_1'], base_n_filters*4, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 55 | |
| 56 | net['upscale3'] = batch_norm(Deconv2DLayer(net['expand_2_2'], base_n_filters*4, 2, 2, crop="valid", nonlinearity=nonlinearity, W=HeNormal(gain="relu"))) |
| 57 | net['concat3'] = ConcatLayer([net['upscale3'], net['contr_2_2']], cropping=(None, None, "center", "center")) |
| 58 | net['expand_3_1'] = batch_norm(ConvLayer(net['concat3'], base_n_filters*2, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 59 | net['expand_3_2'] = batch_norm(ConvLayer(net['expand_3_1'], base_n_filters*2, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 60 | |
| 61 | net['upscale4'] = batch_norm(Deconv2DLayer(net['expand_3_2'], base_n_filters*2, 2, 2, crop="valid", nonlinearity=nonlinearity, W=HeNormal(gain="relu"))) |
| 62 | net['concat4'] = ConcatLayer([net['upscale4'], net['contr_1_2']], cropping=(None, None, "center", "center")) |
| 63 | net['expand_4_1'] = batch_norm(ConvLayer(net['concat4'], base_n_filters, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 64 | net['expand_4_2'] = batch_norm(ConvLayer(net['expand_4_1'], base_n_filters, 3, nonlinearity=nonlinearity, pad=pad, W=HeNormal(gain="relu"))) |
| 65 | # end-snippet-upsampling |
| 66 | |
| 67 | # start-snippet-output |
| 68 | net['output_segmentation'] = ConvLayer(net['expand_4_2'], num_output_classes, 1, nonlinearity=None) |
| 69 | net['dimshuffle'] = DimshuffleLayer(net['output_segmentation'], (1, 0, 2, 3)) |
| 70 | net['reshapeSeg'] = ReshapeLayer(net['dimshuffle'], (num_output_classes, -1)) |
| 71 | net['dimshuffle2'] = DimshuffleLayer(net['reshapeSeg'], (1, 0)) |
| 72 | net['output_flattened'] = NonlinearityLayer(net['dimshuffle2'], nonlinearity=lasagne.nonlinearities.softmax) |