Diffeomorphic unsupervised deep learning model for mono- and multi-modality registration

Abstract

Different from image segmentation, developing a deep learning network for image registration is less straightforward because training data cannot be prepared or supervised by humans unless they are trivial (e.g. pre-designed affine transforms). One approach for an unsupervised deep leaning model is to self-train the deformation fields by a network based on a loss function with an image similarity metric and a regularisation term, just with traditional variational methods. Such a function consists in a smoothing constraint on the derivatives and a constraint on the determinant of the transformation in order to obtain a spatially smooth and plausible solution. Although any variational model may be used to work with a deep learning algorithm, the challenge lies in achieving robustness. The proposed algorithm is first trained based on a new and robust variational model and tested on synthetic and real mono-modal images. The results show how it deals with large deformation registration problems and leads to a real time solution with no folding. It is then generalised to multi-modal images. Experiments and comparisons with learning and non-learning models demonstrate that this approach can deliver good performances and simultaneously generate an accurate diffeomorphic transformation.

Introduction

Image registration consists in constructing a reasonable geometrical correspondence between given two or more images of the same object taken at different times or using the same or different devices in order to locate different or complementary information. Applications of image registration include diverse fields such as astronomy, optics, biology, chemistry, remote sensing and particularly in medical imaging. For an overview of image registration methodology, approaches and applications, we refer to Fischer and Modersitzki ¹ ; Gigengack et al. ² ; Modersitzki ³ ; Oliveira et al. ⁴ ; Sotiras. ⁵ Though the topic is actively studied and useful models exist, there remain many challenges to be tackled mathematically, particularly in registration of images from different modalities. There exist various deformable variational models for image registration where the unknown displacement field u is sought in a properly chosen functional space.^6–11 Generally speaking, the variational problem consists in solving the optimisation problem

$$\underset{\mathbf{u}}{\min }\{\mathcal{L}(\mathbf{u})=S(\mathbf{u})+\frac{\lambda }{2}D(T(\phi ),R)\}$$

(1)

where $\phi (\mathit{x})=\mathit{x}+\mathbf{u}(\mathit{x})$ . In (1), $S(\mathbf{u})$ is a regularisation term which controls the smoothness of u and reflects our expectations by penalising unlikely transformations. The second part $D(T(\phi ),R)$ is a similarity term which measures the goodness of the registration. These models are called non-learning based models as the optimisation problem (1) should be solved for each pair for images T and R. Although various non-learning-based models have been proposed in the recent years and many numerical and computational algorithms have been developed to accelerate the numerical resolution of these models, it remains a very challenging question of achieving both an accurate solution and fast speed for real time applications.

In recent years, deep learning approaches were proposed where the aim is to optimise and learn spatial transformations between pairs of images to be registered,^12–16 often, they require ground-truth deformation fields for the training task. They are called supervised models and their main drawback is the inability to predict transformations that may not be in the same range or class of the training transformations. As example, a deep learning model, which is learned and trained on a dataset where the ground-truth contains only small displacement fields, fails to predict and to give accurate results for large displacement.

In order to remedy these drawbacks, another class of deep leaning models was proposed. These unsupervised models do not require ground-truth deformation fields for training. The deformation fields are self-trained and driven by image similarity metrics computed on the input data. In Jaderberg, ¹² a spatial transformer network is developed to learn transformations for 2D images; however only affine and thin plate spline transformations were used. More general non-parametric transformations were considered in Haskins et al. ¹³ ; Li and Fan ¹⁴ ; Theljani and Chen ^6,25 for mono-modal images. In these approaches, the transformations are controlled by penalising the derivatives of the deformation u in the loss function, which promotes smoothness of predicted transformations. However, this does not guarantee some physicality desired geometric properties of the deformation such as the topology preservation and invertibility. Thus, folding problem can occur in the results which is in practise inappropriate for real life problems where deformation is large and folding can occur. In terms of application, these approaches were only trained and tested on mono-modal images.

To overcome the folding problem, some deep learning approaches have addressed the question of getting diffeomorphic transformations, i.e. topology-preserving and invertible. These conditions were guaranteed by either adding an extra layer to the networks to enforce the output deformation to be diffeomorphic as in Krebs et al., ¹⁷ or by controlling the back-and-forth registration as in Kuang. ¹⁸ In this case, the network is learning a deformation that maps T to R and its possible inverse deformation that maps R to T. Moreover, these approaches were only applied for mono-modal images

In this paper, we propose an alternative approach to overcome the folding problem. The model is unsupervised, i.e. do not require any ground truth data, and is applied for registering mono- and multi-modal images. It can deliver diffeomorphic transformations without adding any extra layer to the neural network or enforcing it compute a deformation and its inverse. In fact, the diffeomorphisms are guaranteed by using a suitable loss function for training that controls the folding in the deformations.

A learning model

In this section, we introduce a learning based diffeomorphic model for both mono- and multi-modal image registration. The idea is that the deformation fields are self-trained by minimizing a specific loss function that guarantees the transformation to be physical. In Figure 1, we present an overview of the method:

Figure 1.

The work-flow of the registration model. The fixed and moving images T and R and first concatenated and fed to the registration network N(T, R). The latter composed of many convolution layers, followed by activation functions, and ended by a final layer that produces the deformation ${\mathbf{u}}_{\theta }$ . The moving image T and the obtained deformation u are passed to an interpolation layer that warps T, i.e. compute $T({\phi }_{\theta })$ . After getting the warped image, it is used to compute the loss function $\mathcal{L}(·)$ and then backpropagate in order to minimise it.

• The pair of images T and R are concatenated and fed to the registration network N(T, R) as a single multichannel image. The registration network, with weight parameters θ, processes the two images through a series of convolutional and pooling layers, and outputs a 2-channel map representing the 2D deformation field, denoted by ${\mathbf{u}}_{\theta }$ . Any network that can capture the image features may work well in our model. Here, we used a light version of the U-Net network with a ReLU activation at the end of each block in the network. The last block is responsible for generating the deformation ${\mathbf{u}}_{\theta }$ .

• Once the estimated transformation ${\mathbf{u}}_{\theta }$ is computed based on the features extracted by the the registration network N(T, R). It is passed to a second component which is a differentiable warping module called Interpolator, that uses ${\mathbf{u}}_{\theta }$ to warp the image T, producing a warped image $T({\phi }_{\theta })$ , where ${\phi }_{\theta }(\mathit{x})=\mathit{x}+{\mathbf{u}}_{\theta }(\mathit{x})$ .

• During training, the parameters θ of the registration network N(T, R) are adjusted by minimising a loss function $\mathcal{L}(·)$ which has form of the energy (1) i.e.

  $$\underset{\theta }{\min }\{\mathcal{L}({\mathbf{u}}_{\theta })=S({\mathbf{u}}_{\theta })+\frac{\lambda }{2}D(T({\phi }_{\theta }),R)\}$$

  (2)

Generally, only the similarity measure $D(·)$ depend on the image modality, whereas same regularisation $S(·)$ could be used for both mono- and multi-modal images. In the sequel, we discuss the choice of the regulariser and the similarity measures.

Regularisation $S(\mathbf{u})$

The regularisation that we use consists in a smoothing constraint on the first- and second order derivatives and a constraint on the determinant of the transformation in order to obtain spatially smooth and plausible solutions Droske and Rumpf ¹⁹ ; Burger et al. ⁷ The second-order derivatives allows getting smooth transformations and penalise affine linear transformations which are not included in the kernel of the first-order derivatives based regularises. ²⁰ The term depending on $\det (\nabla {\phi }_{\theta }(x))$ ensures the map to be locally invertible and then help avoiding the mesh folding problem. More precisely, we consider

$$S({\mathbf{u}}_{\theta })=\left|\right|\nabla {\mathbf{u}}_{\theta }|{|}_2^2+|\left|{\nabla }^2{\mathbf{u}}_{\theta }\right|{|}_2^2+\left|\right|\varphi (\det (\nabla {\phi }_{\theta })||{|}_2^2$$

(3)

where $\varphi (v)=\frac{{(v-1)}^4}{v^2}$ . This term is originally used in the non-learning hyper-elastic model and is known to be very efficient in getting diffeomorphic maps, see Burger et al. ⁷

The similarity measure $D(·)$ depends on the image modality. In the sequel, we discuss the choice of $D(·)$ for each image modality.

Measuring similarity for mono-modal images

Various similarity terms can be used to measure the goodness of the registration. In this work, we use an alternative measure to the correlation coefficient and which is well suitable for measuring linear dependence between images, hence mono-modal images. More precisely, we set

$$D(T({\phi }_{\theta }),R)=\mathbf{CLM}(T({\phi }_{\theta }),R)$$

where CLM is the following combined correlation-like measure:

$$\begin{array}{l}\text{}\mathbf{CLM}(X,Y)=\text{}{\displaystyle {\int }_{\Omega }\text{}}{\left(\frac{(X-m_X)}{\sqrt{\mathbf{Var}(X)}}-\frac{(Y-m_Y)}{\sqrt{\mathbf{Var}(Y)}}\right)}^2\text{}\text{}\text{d}x\\ +{(\sqrt{\mathbf{Var}(X)}+\sqrt{\mathbf{Var}(Y)}-\sqrt{\mathbf{Var}(X+Y)})}^2\end{array}$$

(4)

where m_X and $\mathbf{Var}(X)$ are the mean and the variance of X, respectively.

Measuring similarity for multi-modal images

Multi-modal images are often non-linearly correlated and the similarity measure CLM can’t support multi-modal images as it only measures the linear correlation. In the sequel, we assume that the given multi-modal image pair is not random and has certain connections between them e.g. ‘similar’ shapes or edges.

Parallel level set measure

Various similarity terms can be used in the registration of multi-modal images such as mutual information Pluim et al., ²¹ normalised gradient fields Ruhaak et al. ²² normalised gradients fitting Theljani and Chen ⁶ ; Zhang et al. ²³ In this work, we use parallel level sets similarity measure which is well suitable for measuring alignment between the gradients of two images. More precisely, consider

$$D(T({\phi }_{\theta }),R)=\mathbf{PLS}(T({\phi }_{\theta }),R)$$

where

$$\mathbf{PLS}(X,Y)=\left|\right|\psi (|\nabla X{|}_{\mathrm{ϵ}}|\nabla Y{|}_{\mathrm{ϵ}})-\psi (|\nabla X·\nabla Y{|}_{\mathrm{ϵ}})|{|}_2^2$$

(5)

where $|\nabla X{|}_{\mathrm{ϵ}}=\sqrt{|\nabla X{|}^2+{\mathrm{ϵ}}^2}$ is regularised version of $|\nabla X|$ and ϵ is small non-negative constant. Different function $\psi (·)$ can be used as in Ehrhardt et al., ²⁴ however, here we sued $\psi (x)=x$ .

Data and numerical tests

In the numerical validation, we assess the performance of the proposed model that we call $\mathbf{DL}\mathbf{Model}$ in registering mono-and multi-modal images. Our proposed method was implemented with Tensorflow library. We applied Adam with momentum optimisation algorithm to train the models with a learning rate of 0.0001, and set the batch size 6. The model was trained using a NVIDIA GeForce GTX 1050 Ti GPU. The number of epochs was 1200 for mono-modal images and 400 epochs for the multi-modal case. For the training time, it takes approximately 2 hours to train the model for mono-modal images and 56 minutes for the multi-modal images.

Registration network

We have tested two networks for the registration task. The first one, in Figure 2, is a small network (SN) which contains four blocks containing different kernels with different sizes and ReLU activation is used at the end of each block. The second one is a light version of U-net network, denoted by LU-net ²⁷ . The last block of each network is responsible for generating the deformation ${\mathbf{u}}_{\theta }$ . We illustrate the architecture of this network in Figure 3.

Figure 2.

The architecture of the SN network for registration. The image T and R are concatenated to form a 2-channel image that will be fed to the network. The later output the deformation $\mathbf{u}=(u_1,u_2)$ .

Choice of λ in (2)

We have tested three different values of $\lambda =10,60,200$ . In Figure 4, we display a registration results for a pair of MRI images for these values. We also show the transformed grids and the value $Q_{\min }=\min \det (\nabla {\phi }_{\theta })$ for all values. It is clear that the regularisation parameter λ affects the registration and the mesh qualities. We note that taking large value of λ makes it difficult to get a physical solution, i.e. no folding In fact, for λ = 200 we clearly get a good alignment with $\mathbf{MI}=1.55$ , but this leads to a folded mesh as it can be seen in Figure 5(e), with $Q_{\min }=-0.02$ . However, by decreasing the value of λ, the solution is getting more physically correct with no mesh folding (see the zooms in Figure 5), but the registration quality is not as good as the case where λ is large, i.e. 200.

Figure 3.

(a) The reference image R. (b) The moving image T. (c) Registered image $T({\phi }_{\theta })$ for λ = 200, $\mathbf{MI}=1.55$ . (d) Registered image for λ = 60, MI = 1.49. (e) Registered image for λ = 10, MI= 1.02.

Figure 4.

The transformed meshes for different values of λ for the of images in Figure 3. (a) λ = 200, $Q_{\mathit{min}}=-0.02$ . (b) λ = 60, Q_min = 0.13. (c) λ = 10, Q_min = 0.24. (e) Zoom from the mesh in (a). (f) Zoom from the mesh in (b). Zoom from the mesh in (c). Clearly, a mesh folding occurs for λ = 200.

Figure 5.

The architecture of the light U-net network for registration.

In the rest of the numerical examples, we consider λ = 60.

Part 1: Synthetic data

In the part of the numerical examples, we assess the importance of the term $\left|\right|\varphi (\det (\nabla {\phi }_{\theta })|{|}_2^2$ in the used loss function, and its importance in getting large diffeomorphic deformations by comparing between different common regularisers. We consider 10 synthetic images drawn in order to compare the regularisers for possible large deformations, where each image can serve as a template and reference in the same time (so that we have $2\times {10}^2=200$ pairs of images), see Figure 6. Then, we train and test three different models on this same data, where the losses in these models have the same similarity measures CLM, the same regularisation parameter λ = 60, but different regularisations terms $S(·)$ .

New: It is our proposed model and where

$$S({\mathbf{u}}_{\theta })=\left|\right|\nabla {\mathbf{u}}_{\theta }|{|}_2^2+|\left|{\nabla }^2{\mathbf{u}}_{\theta }\right|{|}_2^2+\left|\right|\varphi (\det (\nabla {\phi }_{\theta })||{|}_2^2.$$

(6)

Total variation: The regularisation in this model is given by $S({\mathbf{u}}_{\theta })=\left|\right|\nabla {\mathbf{u}}_{\theta }|{|}_1$

Diffusion (Diffu): We consider $S({\mathbf{u}}_{\theta })=\left|\right|\nabla {\mathbf{u}}_{\theta }|{|}_2^2$ as regularisation term.

As shown in Figure 7, different $S({\mathbf{u}}_{\theta })$ lead to different results that we can distinguish between them visually and by also checking the value of ${\mathcal{Q}}_{\min }=\min \det (\nabla {\phi }_{\theta })$ . Only the new model gives diffeomorphic deformations since ${\mathcal{Q}}_{\mathrm{min2}}>0$ in this case, i.e. no mesh folding.

Part 2: Real data

We test the performance of the proposed $\mathbf{DL}\mathbf{model}$ in registering real mono- and multi-modal images. Registration quality is evaluated using the mutual information (the larger the better) between the two images $T({\phi }_{\theta })$ and R. We also assess if the map ${\phi }_{\theta }$ is diffeomorphic by checking the minimum of the Jacobian determinant $\det (\nabla {\phi }_{\theta })$ i.e. yes if positive. We tested and compared between the two network architectures SN and LU-net (see Table 2).

We also compare with traditional deformable registration model Burger et al., ⁷ which is a non-learning model and that we call $\mathbf{NL}\mathbf{model}$ . It consists in solving the following optimisation problem

$$\begin{array}{l}\underset{\mathbf{u}}{\min }\left\{\right||\nabla \mathbf{u}|{|}_2^2+\left|\right|{\nabla }^2\mathbf{u}|{|}_2^2+||\varphi (\det (\nabla \phi )\left|\right|{|}_2^2+\frac{\lambda }{2}D(T(\phi ),R)\}\hfill \end{array}$$

(7)

where $D(T({\phi }_{\theta }),R)=\mathbf{CLM}(T({\phi }_{\theta }),R)$ for mono-modal images and $D(T({\phi }_{\theta }),R)=\mathbf{PLS}(T({\phi }_{\theta }),R)$ for multi-modal images. In this model, we solve an optimisation problem w.r.t u for each pair of images T, R. We used the Matlab implementation of FAIR’s Modersitzki ³ for the regularization and we have just implemented the similarity measures CLM and PLS under the framework of FAIR’s package.

Mono-modal images

We trained our model on 160 mono-modal MRI heart images and compare with the classical variational model $\mathbf{NL}\mathbf{model}$ where we solve the optimisation problem (7). We test our model on 20 images and we display 4 comparison tests with the $\mathbf{NL}\mathbf{model}$ in term of runtime and accuracy, see Figure 10. $\mathbf{DL}\mathbf{Model}$ is by far faster and predict the transformation in 1 second for a pair of images with a resolution of 192 × 192. The $\mathbf{NL}\mathbf{model}$ achieves the same result in term of accuracy but it takes more than 1 minute. The comparison between both models is summarised in Table 1.

Table 1.

Comparison between learning DL Model and non-learning NL Model for the images in Figure 10 in Speed (time) and Quality (MI). One sees that DL Model is about 100 times faster for a similar result.

	Examples
	Exp 1	Exp 2	Exp 3	Exp 4
Time (s) for $\mathbf{DL}\mathbf{model}$	1.18 ± 0.6
Time (s) for $\mathbf{NL}\mathbf{model}$	112.54	120.31	107.11	112.5
MI for $\mathbf{DL}\mathbf{model}$	1.49	1.66	1.53	1.57
MI for $\mathbf{NL}\mathbf{model}$	1.52	1.64	1.54	1.59
${\mathbf{Q}}_{\mathit{min}}$ for $\mathbf{DL}\mathbf{model}$	0.45	0.63	0.65	0.67
${\mathbf{Q}}_{\mathit{min}}$ for $\mathbf{NL}\mathbf{model}$	0.61	0.47	0.69	0.54

Figure 6.

Dataset of 10 images used for the training, testing and comparing between different regularisers for large displacement.

Figure 7.

Comparison between different regularisers in registering an unseen pair of images. From left to right, Template T, Reference R, registered using different regularisation terms: i) the diffusion $({\mathbf{Q}}_{\mathit{min}}=0.38)$ , ii) Total-Variation $({\mathbf{Q}}_{\mathit{min}}=-0.24)$ and iii) $S(\mathbf{u})$ in (3) $({\mathbf{Q}}_{\mathit{min}}=0.71)$ , respectively. Clearly iii) gives the best result.

Table 2.

Comparison between the LU-net and SN networks in registering 4 mono-modal MRI heart images.

	Examples
	Exp 1	Exp 2	Exp 3	Exp 4
MI for LU-net	1.53	1.68	1.52	1.62
MI for SN	1.49	1.66	1.53	1.57

We also compared the two used networks in term of registration quality. The comparison details are reported in Table 2. The training curves of the registration model using the two networks are illustrated in Figure 8 and Figure 9. These curves present the learning rate values at each epoch, and we see that the model converges for the two networks after 800 epochs.

Figure 8.

The LU-net: Training and testing losses for the registration model for mono-modal (top) and multi-modal (bottom) images.

Figure 9.

The SN network: Training and testing losses for the registration model for mono-modal (top) and multi-modal (bottom) images.

Figure 10.

Pairwise registration results for 4 pairs of MRI images. Each row represents the result for a pair test. (a) Moving images T. (b) Fixed images R. (c) and (d) are the registered images using $\mathbf{DL}\mathbf{Model}$ and $\mathbf{NL}\mathbf{Model}$ , respectively. The mutual information errors and the values of ${\mathbf{Q}}_{\mathit{min}}$ related to these tests are given in Table 1.

Figure 11.

Pairwise registration results for 4 pairs of MRI-CT images. Each row represents the result for a pair test. (a) Moving images T. (b) Fixed images R. (c) registered images using $\mathbf{DL}\mathbf{Model}$ . (d) and (e) Fused images before and after the registration. The mutual information errors and the values of ${\mathbf{Q}}_{\mathit{min}}$ related to these tests are given in Table 3.

Multi-modal images

For multi-modal images, we trained the network on 120 pairs of CT and MRI images. In Figure 11, we display the prediction results for registering 4 pairs of MRI and CT images. To see the quality visually, we show the fused CT and MRI images for the four examples before and after registration. Clearly after the registration the images are well aligned. We give the mutual information errors and the run-time during the prediction in Table 3.

Table 3.

$\mathbf{DL}\mathbf{model}$ : The Speed (time), Quality (MI) and Q_min vlaues for MR-CT images in Figure 11.

	MR-CT images
	Pair 1	Pair 2	Pair 3	Pair 4
Time (s) for $\mathbf{DL}\mathbf{model}$	0.045 ± 0.002
MI for $\mathbf{DL}\mathbf{model}$	1.95	1.93	1.94	1.94
${\mathbf{Q}}_{\mathit{min}}$ for $\mathbf{DL}\mathbf{model}$	0.97	0.98	0.97	0.97

Conclusions

We have developed and presented an unsupervised deep learning approach for mono-and multi-modal images registration. We tested and compared different choices of regularisation constraints on the deformation fields. The results have shown that control on the Jacobian determinant of the deformation is necessary in the loss function in order to get a diffeomorphic map, mainly for large displacements. The learning model was first designed and tested for mono-modal images. The same learning approach works effectively for multi-modal images, by only changing the similarity measure to fit with the multi-modal setting. For for mono- and multi-modal images, we tested 2 choices of networks. We found that LU-net is recommended.

Future work will consider generalisations to 3 dimensional images registration.