def simple_cnn(input_shape=(32, 32, 2)):
    "Creates a 2-D convolutional encoder-decoder network.\n\n    Parameters\n    ----------\n    input_shape : sequence of ints, optional\n        Input data shape of the form (H, W, C). Default is (32, 32, 2).\n\n    Returns\n    -------\n    model\n        An instance of Keras' Model class.\n\n    Notes\n    -----\n    Given a concatenated pair of static and moving images as input, the\n    CNN computes a dense displacement field that is used to warp the\n    moving image to match with the static image.\n\n    The number of channels in the output (displacement field) is equal\n    to the dimensionality of the input data. For 3-D volumes, it is 3,\n    and for 2-D images, it is 2. The first channel comprises\n    displacement in the x-direction and the second comprises\n    displacement in the y-direction.\n    "
    out_channels = 2
    inputs = layers.Input(shape=input_shape)
    x = layers.Conv2D(32, kernel_size=3, strides=2, padding='same', activation='relu')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(32, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPool2D(pool_size=2)(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPool2D(pool_size=2)(x)
    x = layers.Conv2D(128, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(64, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(32, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(32, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(16, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(16, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(out_channels, kernel_size=1, strides=1, padding='same')(x)
    model = tf.keras.Model(inputs, x, name='simple_cnn')
    return model
    '\n    Differntiable image sampling\n    References:\n    1.   https://github.com/tensorflow/models/blob/master/research/transformer/spatial_transformer.py\n    2.   Jaderberg, Max, Karen Simonyan, and Andrew Zisserman. "Spatial\n        transformer networks." Advances in neural information processing\n        systems. 2015. https://arxiv.org/pdf/1506.02025.pdf\n    3.   *Spatial* Transformer Networks by Kushagra Bhatnagar https://link.medium.com/0b2OrmqVO5\n    '
Creates a 2-D convolutional encoder-decoder network.

Parameters
----------
input_shape : sequence of ints, optional
    Input data shape of the form (H, W, C). Default is (32, 32, 2).

Returns
-------
model
    An instance of Keras' Model class.

Notes
-----
Given a concatenated pair of static and moving images as input, the
CNN computes a dense displacement field that is used to warp the
moving image to match with the static image.

The number of channels in the output (displacement field) is equal
to the dimensionality of the input data. For 3-D volumes, it is 3,
and for 2-D images, it is 2. The first channel comprises
displacement in the x-direction and the second comprises
displacement in the y-direction.
def simple_cnn(input_shape=(32, 32, 2)):
    "Creates a 2-D convolutional encoder-decoder network.\n\n    Parameters\n    ----------\n    input_shape : sequence of ints, optional\n        Input data shape of the form (H, W, C). Default is (32, 32, 2).\n\n    Returns\n    -------\n    model\n        An instance of Keras' Model class.\n\n    Notes\n    -----\n    Given a concatenated pair of static and moving images as input, the\n    CNN computes a dense displacement field that is used to warp the\n    moving image to match with the static image.\n\n    The number of channels in the output (displacement field) is equal\n    to the dimensionality of the input data. For 3-D volumes, it is 3,\n    and for 2-D images, it is 2. The first channel comprises\n    displacement in the x-direction and the second comprises\n    displacement in the y-direction.\n    "
    out_channels = 2
    inputs = layers.Input(shape=input_shape)
    x = layers.Conv2D(32, kernel_size=3, strides=2, padding='same', activation='relu')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(32, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPool2D(pool_size=2)(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPool2D(pool_size=2)(x)
    x = layers.Conv2D(128, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(64, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(64, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(32, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(32, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2DTranspose(16, kernel_size=2, strides=2, padding='same')(x)
    x = layers.Conv2D(16, kernel_size=3, strides=1, padding='same', activation='relu')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(out_channels, kernel_size=1, strides=1, padding='same')(x)
    model = tf.keras.Model(inputs, x, name='simple_cnn')
    return model
    '\n    Differntiable image sampling\n    References:\n    1.   https://github.com/tensorflow/models/blob/master/research/transformer/spatial_transformer.py\n    2.   Jaderberg, Max, Karen Simonyan, and Andrew Zisserman. "Spatial\n        transformer networks." Advances in neural information processing\n        systems. 2015. https://arxiv.org/pdf/1506.02025.pdf\n    3.   *Spatial* Transformer Networks by Kushagra Bhatnagar https://link.medium.com/0b2OrmqVO5\n    '