Source code for

# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

import numpy as np
import tensorflow as tf
from tensorflow.contrib.layers.python.layers import regularizers

from niftynet.engine.application_initializer import GlorotUniform
from niftynet.layer.convolution import ConvolutionalLayer as Conv
from niftynet.layer.downsample_res_block import DownBlock as DownRes
from niftynet.layer.fully_connected import FullyConnectedLayer as FC
from niftynet.layer.grid_warper import AffineGridWarperLayer as Grid
from niftynet.layer.layer_util import infer_spatial_rank
from niftynet.layer.linear_resize import LinearResizeLayer as Resize
from import BaseNet

[docs]class INetAffine(BaseNet): """ ### Description This network estimates affine transformations from a pair of moving and fixed image: Hu et al., Label-driven weakly-supervised learning for multimodal deformable image registration, arXiv:1711.01666 Hu et al., Weakly-Supervised Convolutional Neural Networks for Multimodal Image Registration, Medical Image Analysis (2018) see also: ### Building blocks [DOWN CONV] - Convolutional layer + Residual Unit + Max pooling [CONV] - Convolutional layer [FC] - Fully connected layer, outputs the affine matrix [WARPER] - Grid resampling with the obtained affine matrix ### Diagram INPUT PAIR --> [DOWN CONV]x4 --> [CONV] --> [FC] --> [WARPER] --> DISPLACEMENT FIELD ### Constraints - input spatial rank should be either 2 or 3 (2D or 3D images only) """
[docs] def __init__(self, decay=1e-6, affine_w_initializer=None, affine_b_initializer=None, acti_func='relu', name='inet-affine'): """ :param decay: float, regularisation decay :param affine_w_initializer: weight initialisation for affine registration network :param affine_b_initializer: bias initialisation for affine registration network :param acti_func: activation function to use :param name: layer name """ BaseNet.__init__(self, name=name) self.fea = [4, 8, 16, 32, 64] self.k_conv = 3 self.affine_w_initializer = affine_w_initializer self.affine_b_initializer = affine_b_initializer self.res_param = { 'w_initializer': GlorotUniform.get_instance(''), 'w_regularizer': regularizers.l2_regularizer(decay), 'acti_func': acti_func} self.affine_param = { 'w_regularizer': regularizers.l2_regularizer(decay), 'b_regularizer': None}
[docs] def layer_op(self, fixed_image, moving_image, is_training=True, **unused_kwargs): """ :param fixed_image: tensor, fixed image for registration (defines reference space) :param moving_image: tensor, moving image to be registered to fixed :param is_training: boolean, True if network is in training mode :return: displacement fields transformed by estimating affine """ spatial_rank = infer_spatial_rank(moving_image) spatial_shape = fixed_image.get_shape().as_list()[1:-1] # resize the moving image to match the fixed moving_image = Resize(spatial_shape)(moving_image) img = tf.concat([moving_image, fixed_image], axis=-1) res_1 = DownRes(self.fea[0], kernel_size=7, **self.res_param)(img, is_training)[0] res_2 = DownRes(self.fea[1], **self.res_param)(res_1, is_training)[0] res_3 = DownRes(self.fea[2], **self.res_param)(res_2, is_training)[0] res_4 = DownRes(self.fea[3], **self.res_param)(res_3, is_training)[0] conv_5 = Conv(n_output_chns=self.fea[4], kernel_size=self.k_conv, with_bias=False, feature_normalization='batch', **self.res_param)(res_4, is_training) if spatial_rank == 2: affine_size = 6 elif spatial_rank == 3: affine_size = 12 else: tf.logging.fatal('Not supported spatial rank') raise NotImplementedError if self.affine_w_initializer is None: self.affine_w_initializer = init_affine_w() if self.affine_b_initializer is None: self.affine_b_initializer = init_affine_b(spatial_rank) affine = FC(n_output_chns=affine_size, feature_normalization=None, w_initializer=self.affine_w_initializer, b_initializer=self.affine_b_initializer, **self.affine_param)(conv_5) grid_global = Grid(source_shape=spatial_shape, output_shape=spatial_shape)(affine) return grid_global
[docs]def init_affine_w(std=1e-8): """ :param std: float, standard deviation of normal distribution for weight initialisation :return: random weight initialisation from normal distribution with zero mean """ return tf.random_normal_initializer(0, std)
[docs]def init_affine_b(spatial_rank, initial_bias=0.0): """ :param spatial_rank: int, rank of inputs (either 2D or 3D) :param initial_bias: float, initial bias :return: bias initialisation for the affine matrix """ if spatial_rank == 2: identity = np.array([[1., 0., 0.], [0., 1., 0.]]).flatten() elif spatial_rank == 3: identity = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0]]).flatten() else: tf.logging.fatal('Not supported spatial rank') raise NotImplementedError identity = identity.reshape([1, -1]) identity = np.tile(identity, [1, 1]) return tf.constant_initializer(identity + initial_bias)