Unsupervised Multi-Modal Image Registration via Geometry Preserving Image-to-Image Translation

Our registration network ($R=(R_{\Phi},R_{S})$) is a spatial transformation network (STN) composed of a fully-convolutional network $R_{\Phi}$ and a re-sampler layer $R_{S}$. The transformation we apply is a non-linear dense deformation - allowing non-uniform mapping between the images and hence gives accurate results. Next we give an in-depth description about each component.

$\mathbf{R_{\Phi}}$ - Deformation Field Generator: The network takes two input images, $I_{a}$ and $I_{b}$, and produces a deformation field $\phi=R(I_{a},I_{b})$ describing how to non-rigidly align $I_{a}$ to $I_{b}$. The field is an $H_{A}\times W_{A}$ matrix of $2$-dimensional vectors, indicating the deformation direction for each pixel $(i,j)$ in the input image $I_{a}$.

$\mathbf{R_{S}}$ - Re-sampling Layer: This layer receives the deformation field $\phi$, produced by $R_{\Phi}$, and applies it on a source image $I_{s}$. Here, the source image is not necessarily $I_{a}$ and it could be from either domains - $\mathcal{A}$ or $\mathcal{B}$. Specifically, the value of the transformed image $R_{S}(I_{s},\phi)$ at pixel $\mathbf{v}=(i,j)$ is given by Equation 1:

where $\phi(\mathbf{v})=(\Delta y,\Delta x)$ is the deformation generated by $R_{\Phi}$ at pixel $\mathbf{v}=(i,j)$, in the $x$ and $y$-directions, respectively.

To avoid overly distorting the deformed image $R_{S}(I_{s},\phi)$ we restrict $R_{\Phi}$ from producing non-smooth deformations. We adapt a common regularization term that is used to produce smooth deformations. In particular, the regularization loss will encourage neighboring pixels to have similar deformations. Formally, we seek to have small values of the first order gradients of $\phi$, hence the loss at pixel $\mathbf{v}=(i,j)$ is then given by:

where $N(\mathbf{v})$ is a set of neighbors of the pixel $\mathbf{v}$, and $B\left(\mathbf{u},\mathbf{v}\right)$ is a bilateral filter [32] used to reduce over-smoothing. Let $O_{s}=R_{S}(I_{s},\phi)$ to be the deformed image produced by $R_{S}$ on input $I_{s}$, then the bilateral filter is given by:

There are two important notes about the bilateral filter $B$ in Equation 3. First, the bilateral filtering is with respect to the transformed image $O_{s}$ (at each forward pass), and secondly, the term $B\left(I_{s},\mathbf{u},\mathbf{v}\right)$ is a treated as constant (at each backward pass). The latter is important to avoid $R_{\Phi}$ alternating pixel values so that $B(u,v)\approx 0$ (e.g., it could change pixels so that $\left\lVert O_{s}[\mathbf{u}]-O_{s}[\mathbf{v}]\right\rVert$ is relatively large), while the former allows better exploration of the solution space.

In our experiments we look at the $3\times 3$ neighborhood of $\mathbf{v}$, and set $\alpha=1$. The overall smoothness loss of the network $R$, denoted by $\mathcal{L}_{smooth}\left(R\right)$, is the mean value over all pixels $\mathbf{v}\in\left\{1,\dots,H_{A}\right\}\times\left\{1,\dots,W_{A}\right\}$.

A key challenge of our work is to train the image-to-image translation network $T$ to be geometric preserving. If $T$ is geometric preserving, it implies that it only performs photo-metric mapping, and as a consequence the registration task is performed solely by the registration network $R$. However, during our experiments, we observed that $T$ tends to generate fake images that are spatially aligned with the ground truth image, regardless of $R$’s accuracy.

To avoid this, we could restrict $T$ from performing any spatial alignment by reducing its capacity (number of layers). While we did observe that reducing $T$’s capacity does improve our registration network’s performance, it still limits the registration network from doing all the registration task (See supplementary materials).

To implicitly encourage $T$ to be geometric preserving we require that $T$ and $R$ are commutative, i.e., $T\circ R=R\circ T$. In the following we formally define both $T\circ R$ and $R\circ T$:

To train $R$ and $T$ to generate fake samples that are similar to those in domain $\mathcal{B}$, we use an $L1$-reconstruction loss:

where minimizing the above implies that $T\circ R\approx R\circ T$.

We use conditional GAN (cGAN)[22] as our adversarial loss for training $D$, $T$ and $R$. The objective of the adversarial network $D$ is to discriminate between real and fake samples, while $T$ and $R$ are jointly trained to fool the discriminator. The cGAN loss for $T\circ R$ and $R\circ T$ is formulated below:

The total objective is given by:

where we are opt to find $T^{*}$ and $R^{*}$ such that $T^{*},R^{*}=\underset{R,T}{\arg\min}\mathcal{L}(T,R)$. Furthermore, in our experiments, we set $\lambda_{R}=100$ and $\lambda_{S}=200$.

Our code is implemented using PyTorch 1.1.0 [24] and is based on the framework and implementation of Pix2Pix [13], CycleGAN [36] and BiCycleGAN [37]. The network $T$ is an encoder-decoder network with residual connections [1], which is adapted from the implementation in [15]. The registration network is U-NET based [26] with residual connections in the encoder-path and the output path. In all residual connections, we use Instance Normalization Layer [33]. All networks were initialized by the Kaiming [11] initialization method. Full details about the architectures is provided in the supplementary material.

The experiments were conducted on single GeForce RTX 2080 Ti. We use Adam Optimizer [17] on a mini-batch of size 12 with parameters $lr=1\times e^{-4}$, $\beta_{1}=0.5$ and $\beta_{2}=0.999$. We train our model for 200 epochs, and activate linear learning rate decay after 100 epochs.

Registration Accuracy Metric. We manually annotated 100 random pairs of test images. We tagged 10-15 pairs of point landmarks on the source and target images which are notable and expected to match in the registration process (See Figure 3). Given a pair of test images, $I_{a}$ and $I_{b}$, with a set of tagged pairs ${(x_{a_{i}},y_{a_{i}}),(x_{b_{i}},y_{b_{i}})}_{i\in[N]}$, denote $(x^{\prime}_{a_{i}},y^{\prime}_{a_{i}})$ as deformed sample points $(x_{a_{i}},y_{a_{i}})$. The accuracy of the registration network $R$ is simply the average Euclidean distance between the target points and the deformed source points:

Furthermore, we used two type of annotations. The first type of annotation is located over salient objects in the scene (the blue points in Figure 3). This is important because in most cases, down-stream tasks are affected mainly by the alignment of the main object in the scene in both modalities. The second annotation is performed by picking landmark points from all objects across the scene.

Quantitative Evaluation. Due to limited access to source code and datasets of related works, we conduct several experiments that demonstrate the power of our method with respect to different aspects of previous works. As the crux of our work is the alleviation of the need for cross-modality similarity measures, we trained our network with commonly used cross-modality measures. In Table 1 we show the registration accuracy of our registration network $R$ when trained with different loss terms. Specifically, we used Normalized Cross Correlation as it is frequently used in unsupervised multi-modal registration methods. Furthermore, we trained our network with Structural Similarity Index Metric (SSIM) on edges detected by Canny edge-detector [4] from both the deformed and target image. Finally, we attempt to train our registration network $R$ by maximizing the Normalized Cross Correlation between the edges of the deformed and the target image. As can be seen from Table 1, training the registration network $R$ using prescribed cross-modality similarity measures do not perform well. Further, using these $NCC$ produces noisy results, while using SSIM gives smooth but less accurate registration (see supplemental materials).

Furthermore, we also tried using traditional descriptors such as SIFT [19] in order to match corresponding key points from the source and target image. We use these key points to register the source and target images by estimating the transformation parameters between them. However, these descriptors are not designed for multi-modal data, and hence they fail badly to be used on our dataset.

Instead, we train a CycleGAN [36] network to translate between the two modalities at hand, without any supervision to match the ground truth. CycleGAN, like other unsupervised image-to-image translation networks, is not trained to generate images matching ground truth samples, thus, geometric transformation is not explicitly required from the translation network. Once trained, we use one of the generators in the CycleGAN, the one that maps between domain $\mathcal{A}$ to domain $\mathcal{B}$ to translate the input image $I_{a}$ onto modality $\mathcal{B}$. Assuming this generator is both geometry preserving and translates well between the modalities, it is expected that it also match well between feature of the fake sample and the target image. Thus, we extracted SIFT descriptors from the generated images by the CycleGAN translation network, and extracted SIFT features from the target image $I_{b}$. We then matched these features and estimated the needed spatial registration. The registration accuracy using this method is significantly better than directly using SIFT [19] features on the input image $I_{a}$. The results are shown in Table 1. Further visual results and details demonstrating this method are provided in the supplementary material.

Qualitative Evaluation.

Figure 5 shows that our registration network successfully aligns images from different pair of modalities and handles different alignment cases. For example, the banana leaves in the first raw in Figure 5(a) are well-aligned in the two modalities. Our registration network maintains this alignment and only deforms the background for full alignment between the images. This can be seen from the deformation field visualization [2], where little deformation is applied on the banana plant, while most of the deformation is applied on the background. Furthermore, in the last row in Figure 5(a), there is little depth variation in the scene because the banana plant is small, hence a uniform deformation is applied across the entire image. To help measuring the alignment success, we overlay (with semi-transparency) the plant in image B on top of both image A before and after the registration. This means that the silhouette has the same spatial location in all images (the original image B, image A before and after the registration). Lastly, we achieve similar success in the registration between RGB and IR images (see Figure 5(b)).

It is worth mentioning that in some cases, the deformation field points to regions outside the source image. In those cases, we simply sample zero values. This happens because the the target image content (i.e., $I_{b}$) in these regions is not available in the source image (i.e., $I_{a}$). We provide more qualitative results in the supplemental materials.

Next, we present a series of ablation studies that analyze the effectiveness of different aspects in our work. First, we show that training both compositions (i.e our presented two training flows) of $T$ and $R$ indeed encourages a geometric preserving translator $T$. Additionally, we analyze the impact of the different loss terms on the registration network’s accuracy. We further show the effectiveness of the bilateral filtering, and that it indeed improves the registration accuracy. All experiments, unless otherwise stated, were conducted without the bilateral filtering.

Geometric-Preserving Translation Network.

To evaluate the impact of training of $T$ and $R$ simultaneously with the two training flows proposed in Figure 2, we compare the registration accuracy of our method with that of training models with either $T\circ R$ or $R\circ T$. As can be seen from Figure 6, training both combinations yields a substantial improvement in the registration accuracy (shown in blue), compared to each training flow (i.e., $T\circ R$ and $R\circ T$) separately. Moreover, while the reconstruction loss of $T\circ R$ (shown in read) is lowest among the three options, it does not necessarily indicate a better registration. This is because in this setting the translation network $T$ implicitly performs both the alignment and translation tasks. Conversely, when training with $R\circ T$ only (shown in green), the network $R$ is unstable and at some point it starts to alternate pixel values, essentially taking on the role of a translation network. Since $R$ is only geometry-aware by design it fails to generate good samples. This is indicated by how fast the discriminator detects that the generated samples are fake (i.e., the adversarial loss decays fast). Visual results are provided in the supplementary materials.

Loss ablation. It has been shown in previous works [37, 13, 36] that training an image-to-image translation network with both a reconstruction and an adversarial loss yields better results. In particular, the reconstruction loss stabilizes the training process and improves the vividness of the output images, while the adversarial loss encourages the generation of samples matching the real-data distribution.

The main objective of our work is the production of a registration network. Therefore, we seek to understand the impact of both losses (reconstruction and adversarial) on the registration network. To understand the impact of each loss, we train our model with different settings: each time we fix either $R$ or $T$’s weights with respect to one of the loss functions. The registration accuracy is presented in Table 2. Please refer to the supplementary material for qualitative results. As can be seen in these figures, training $R$ only with respect to the reconstruction loss leads to overly sharp, but unrealistic images where the deformation field creates noisy artifacts. On the other hand, training $R$ only with respect to the adversarial loss creates realistic images, but with inexact alignment. This is especially evident in Table 2 where training $R$ with respect to the reconstruction loss achieves a significant improvement in the alignment, and the best accuracy is obtained when the loss terms are both used to update all the networks weights.

Bilateral Filtering Effectiveness Using bilateral filtering to weigh the smoothness loss allows us, in effect, to encourage piece-wise smoothness on the deformation map. As can be seen in Table 3, this enhances the precision of the registration. These results suggest that using segmentation maps for controlling the smoothness loss term could be beneficial.