Deformation equivariant cross-modality image synthesis with paired non-aligned training data

Cross-modality medical image synthesis has gained a lot of attention in recent years with multiple proposed clinical applications (Wang et al., 2021b). The conditional GANs-based architecture pix2pix (Isola et al., 2017) is widely used when paired and aligned data are available as it is based on an assumption of pixel-to-pixel correspondence between training images of different modality. Whereas CycleGAN (Zhu et al., 2017) can be used if one does not have paired or aligned data.

Paired training images in the medical context are typically not aligned, and hence for pixel-to-pixel training they are registered into the same coordinate system. However, registration is never perfect and pixel-to-pixel losses are very sensitive to registration errors reducing the synthesis quality especially on difficult to register areas with large internal anatomic motion such as on pelvis area (Wang et al., 2021b).

In pixel-to-pixel setting, different approaches have been proposed to mitigate for the remaining registration errors (Chen et al., 2020; Joyce et al., 2017; Kazemifar et al., 2019; Leynes et al., 2018; Yu et al., 2019). Most similarly to our work, Kong et al. (2021) combine cross-modality image synthesis network with registration network to enable training with non-aligned data. However, we argue that their method does not robustly mitigate for registration errors especially with real world data sets. Additionally they use unconditional adversarial training which has been shown to be inferior to conditioning the discriminator with input images. In this work we aim to solve both of these problems.

While performing worse than pix2pix when paired and aligned data are available, unsupervised CycleGAN is more robust to misalignments due to its cycle consistency loss (Kaji and Kida, 2019; Wang et al., 2021b). However, if the misalignments between the modalities are systematic and severe, CycleGAN can also fail to produce geometrically aligned predictions. Approaches, similar to ones used with pix2pix, have also been employed with CycleGAN (Zhang et al., 2018; Hiasa et al., 2018; Kida et al., 2019). Wang et al. (2018, 2019, 2021a) propose network architectures combining cross-modality image synthesis and registration, together with mutual information loss between the input and the prediction to enforce similar geometry.

Deformable medical image registration using deep learning has also gained popularity recently (Fu et al., 2020). With stationary velocity field parametrization, one can generate diffeomorphic deformations (Arsigny et al., 2006; Ashburner, 2007). This kind of methodology was applied to deep learning by Dalca et al. (2018). Some architectures such as the one by De Vos et al. (2019) combine affine or rigid registration together with a separate deformable registration resulting in multi-stage registration approach.

From cross-modality registration methods, the method by Arar et al. (2020) is closest to our work. They train a cross-modality registration network by simultaneously training a cross-modality image synthesis network which they encourage to be equivariant to deformations predicted by the registration network. This is done by applying the predicted deformation both before and after the image synthesis network and comparing both of them to the label. We instead use simulated deformations for encouraging deformation equivariance which we argue to be more robust approach. Our method of encouraging deformation equivariance is similar to the method by Pielawski et al. (2020) where they train their network for rotational equivariance.

Very recently, Chen et al. (2022) use contrastive learning based loss for enforcing geometric (or shape) similarity of the image synthesis in an otherwise similar setting to Arar et al.

The equivariance similarity loss and the commutation loss require some way to simulate anatomically realistic deformations for calculating them. An ideal set would be the set of all anatomically realistic deformations for each sample, but anatomically realistic non-affine deformations are very difficult to simulate. A natural question is then whether a significantly smaller set of deformations would be enough especially given the inductive biases of the used architectures.

Given that $F$ is a convolutional network it is (roughly) translation equivariant. Hence intuitively simulating only globally affine deformations might be enough since any diffeomorphic deformation is "locally affine". In our experiments globally affine deformations were indeed enough, and we ended up using only rotations and flips. As a result, deformations can be sampled from the same distribution for each sample. Applying affine deformations to images is also computationally efficient and numerically accurate. With more challenging data, scaling, shearing or even elastic deformations might have to be used.

Flips deviate from the idea of the set of deformations being anatomically realistic. However, physically flipping is a relatively well defined action, although in practice usually impossible to execute. Most imaging methods image the physical world in such a way that flipping the imaged subject should also flip the resulting image.

In addition to the losses presented so far, we want to incorporate adversarial loss to the training in order to improve the appearance and also clinical quality of the predictions. In adversarial training, an additional discriminator network is trained to classify whether an image fed to the network is real or fake and can be used for guiding the generator network responsible for generating synthetic images. We want to employ a conditional adversarial training setting, similar to pix2pix (Isola et al., 2017), wherein the input image is also fed to the discriminator. This is different to the approach taken by Kong et al. (2021) where they feed only the prediction or the label. Conditioning the discriminator on input images results in generally better predictions (Isola et al., 2017).

Let now $D$ be the discriminator network receiving an input image as the first argument and either a label or a prediction as the second argument. The conditional adversarial learning objective is defined as follows:

The discriminator is trained to maximize the loss and the generator is trained to minimize it with the training executed in turns while holding the weights of the other network constant. Placeholder texts are used as we are yet to derive the optimal way to feed the data to the discriminator. Misaligned training data will require care in how that is done. To be more precise, the following three points need to be taken into account:

1. Predictions and labels have to be fed in the same coordinate system since the input domain and the label domain might have systematic geometric differences which would encourage the predictions to be misaligned with the inputs.

2. Inputs and their corresponding predictions can not be fed in the exactly same coordinate system since even if the predictions are registered to the labels or vice versa, the labels will not be exactly aligned with the inputs, especially in the beginning of the training. That would encourage misaligned predictions.

3. Interpolation acts as a low-pass filter especially in areas where the image is stretched. As a result, predictions registered to labels can not be directly compared with the labels as the discriminator will learn to notice the missing high frequencies.

Points 1) and 2) would suggest to feed labels and predictions in the label coordinates but 3) would suggest that we can not deform the predictions to the labels either. As a solution we propose to train two separate registration networks: one for cross-modality registration of labels to inputs and one for intra-modality registration of predictions to possibly imperfectly registered labels. The adversarial comparison can then be done between the registered labels and the predictions registered to the registered labels. The proposed approach will solve all the three problems mentioned above: 1) The comparison will be done in the same coordinate system. 2) If the label and the input are imperfectly aligned, the prediction can still be separately registered to the registered label hence removing the incentive for misaligned predictions. 3) As the training progresses, most of the movement should be included in the registration of the labels to the inputs, allowing the predictions registered to the registered labels to contain at least as high frequency information as the registered labels.

Let us now denote the predicted cross-modality deformation (approximately) mapping coordinates of ${\tilde{y}}^{\left(i\right)}$ to $x^{\left(i\right)}$ as $d_{\text{cross}}^{\left(i\right)}$ and the predicted intra-modality deformation (approximately) mapping coordinates of $F\left(x^{\left(i\right)}\right)$ to ${d_{\text{cross}}^{\left(i\right)}}^{\ast }{\tilde{y}}^{\left(i\right)}$ as $d_{\text{intra}}^{\left(i\right)}$. We will look at how these are obtained in Section 4.3.

From the adversarial loss perspective, we want to treat $d_{\text{cross}}^{\left(i\right)}$ as constant since the cross-modality registration would not benefit from the adversarial loss and might even result in unexpected optima. This approach corresponds to the normal GAN training where only the second term is used in updating the generator. The proposed adversarial loss is then

The loss function can be improved even further by employing similar idea to the equivariance similarity loss. To simultaneously prevent discriminator from over-fitting and implicitly enforce equivariance to set of deformations, we propose to further modify the loss to the following form (which we label equivariance adversarial loss):

Here same deformations $t$ can be used which are used for the equivariance similarity loss and the commutation loss. The core idea is to augment the inputs to the discriminator with the deformations $t$ and the formulation follows from it naturally.

As discussed, the registration will be divided into cross-modality registration for registering labels to inputs and intra-modality registration for registering predictions to the registered labels. While the cross-modality registration receives pairs of different modality as inputs, it is trained with intra-modality loss based on the synthesised image $F\left(x^{\left(i\right)}\right)$ similarily to the intra-modality registration.

In principle, the registration networks can predict the deformation in any suitable form. We split the cross-modality registration into rigid registration and elastic registration. The two-stage architecture makes it significantly easier for the model to handle large deformations. For intra-modality registration, we do not use two-stage architecture as cross-modality registration should take care of most of the movement.

We generate elastic deformations from stationary velocity fields to enforce diffeomorphic deformations and to allow inverting the deformations. From a stationary velocity field the final diffeomorphic deformation is obtained by integrating the field over itself over a unit time. In group theory, this can be seen as exponentiation of a member of a lie algebra (Arsigny et al., 2006), and hence we denote the integration by $\exp$. Exponentiation can be estimated efficiently by the scaling and squaring method (Arsigny et al., 2006; Dalca et al., 2018). The velocity fields are predicted in the same resolution as the images.

Throughout the architecture, images are resampled by deformations, but the sampled locations might be outside the image. We connect each image with a mask that can initially represent invalid regions in the image. Each time an image is resampled, the mask is updated with the regions resampled from outside the image. In similarity losses, we then compare images only within intersections of the masks. Masks contain discrete values and no gradients flow through the masks during the backward-pass, preventing the optimization of the masks themselves. The same procedure is done also for the deformations as they are also resampled by other deformations.

Invalid region masks are also fed for the registration networks since for invalid regions only the regularization should affect the generated deformation.

Additionally, the masks of the registered labels and the predictions registered to the registered labels might be systematically different, which the discriminator could use for separating the images. We mitigate for that by multiplying each image fed to the discriminator by the intersection of the masks of the images compared.

The overall loss function can be written as

where $L_{\text{sim}}^{\text{intra}}$ is either $L_{\text{eq-sim}}^{\text{intra}}$ or $L_{\text{def-sim}}^{\text{intra}}$, and $\lambda ,\gamma ,\delta \in R$ are loss function weights. Each loss component affects only the weights of the sub-networks written in the curly braces. Note that $D$ is trained to maximize the loss whereas the other networks are trained to minimize it.

Actual values used for the loss function weights are given in the supplementary materials.

Synthetic "multimodal" data sets were created using images from COCO data set (Lin et al., 2014) with unmodified images as input images. Label images were generated by circularly swapping the RGB color channels of the input images and by deforming them with simulated deformations. All images were centrally cropped to resolution $(400,400)$. Four data set were generated: LR (Large Random), SR (Small Random), LC (Large Constant), and SC (Small Constant). Used deformation parameters are displayed in Table 1. Training, validation, and test sets all contained $4113$ images.

With this data set, all the experiments were conducted without the adversarial loss to study separately the effects of deformation equivariance encouraging losses. The six models trained using each of the data sets are listed in Table 2.

No model was trained with aligned data as it would be easily learned perfectly with data set this easy.

In recent years, a significant amount of research has emerged on applying deep learning to cross-modality brain MRI synthesis. Synthetically generated modalities have many possible down-stream use cases such as segmentation, classification, detection and diagnosis. (Xie et al., 2022)

We used brain images from Information eXtraction from Images (IXI) data set ¹¹1http://brain-development.org/ixi-data set/ to generate a semi-synthetic 3D data set for T2 to PD (proton density) synthesis, like in Wang et al. (2021a). The main reason for choosing this task was that brain T2 and PD images are initially well aligned and hence provide good ground truths. Most of the images were already in resolution $(1.25\text{mm},0.9375\text{mm},0.9375\text{mm})$, the rest were also resampled to the same resolution. All the images were normalized based on brain white-matter using the implementation by Reinhold et al. (2019). We used $192$ images for training, $19$ for validation, and $365$ for testing.

Simulated deformations were added to the PD images. Translations were sampled from range ($3.0\text{mm}$, $10.0\text{mm}$), rotations from range (${1.0}^{\circ }$, ${4.0}^{\circ }$), $\sigma$ from range ($24.0\text{mm}$, $72.0\text{mm}$), and $m$ from range ($-15.0\text{mm}$, $15.0\text{mm}$) with values sampled independently for each dimension or rotation axis. The elastic component mean parameter $\mu$ was sampled uniformly from image coordinates. The distribution is intentionally skewed to make the non-desired outcome of over-learning the deformation already in $F$ more attractive.

The six models trained with this data set are listed in Table 3. Note that the model NoReg + DefCondAdv + Aug is an oracle as it is trained with the aligned data set. It’s performance should provide a good upper boundary on the performance of the other models.

All of the models with this data set were trained by sampling random image patches of size $(64,64,64)$ from the whole training data set. Additionally the inputs were augmented with low-amplitude noise.

Virtual histopathology staining using deep learning has emerged as an active research topic in recent years, and has been primarily driven by GAN-based methods (Bayramoglu et al., 2017; Rana et al., 2020; Rivenson et al., 2019). However, a majority of the methods require elastic registration of inputs and labels. Our method simplifies training of the network by eliminating the need to elastically register image pairs.

We used a non-public data set containing unstained and stained tissue whole slides image (WSI) pairs. These are essentially ultra high resolution gigapixel images, and virtually staining the unstained tissue WSIs is a highly non-trivial task. Pre-clinical murine prostate tissue samples were prepared at the University of Eastern Finland, Kuopio. Material used was surplus tissue from previous studies (Latonen et al., 2017; Valkonen et al., 2017) where all animal experimentation and care procedures were carried out in accordance with guidelines and regulations of the national Animal Experiment Board of Finland, and were approved by the board of laboratory animal work of the State Provincial Offices of South Finland (licence number ESAVI/6271/04.10.03/2011). The tissue samples were first scanned without staining. This was followed by hematoxylin and eosin (H&E) staining of the unstained tissue samples, and then the stained samples were scanned again. The samples were scanned using Thunder Imager 3D Tissue slide scanner (Leica Microsystems, Wetzlar, Germany) equipped with DMC2900 camera at 40X magnification level with a pixel size of 0.353µm. Total of 17 WSI pairs were included in the data set each with resolution of approximately $40\text{k}\times 40\text{k}$ from which $9$ were used for training, $1$ for validation, and $7$ for testing.

Inputs and labels were coarsely registered and the alignment seems superficially good. However, upon a closer inspection clear misalignments are present.

The six models trained with this data set are listed in Table 4. No oracle model was trained as we did not have ground truth registrations for this data set. For evaluation, however, we did additionally register the images using an open source cross-modality whole slide image registration tool called wsireg²²2https://github.com/NHPatterson/wsireg. The WSI pairs were registered in two steps, first rigidly for global alignment and then elastically for more granular correspondence between the modalities.

The models were trained by sampling random image patches of size $512\times 512$ from the whole training data set. Additionally inputs were augmented with low-amplitude noise.

For image synthesis evaluation, we used three metrics: peak-signal-to-noise-ratio (PSNR), structural similarity index (SSIM) (Wang et al., 2004), and normalized mutual information (NMI) (Studholme et al., 1999). NMI was applied between inputs and predictions as opposed to inputs and aligned labels as it is used here for measuring the geometric similarity of the predictions to the corresponding inputs.

Additionally we evaluated the accuracy of the overall predicted deformation with a metric which we label mean deformation error (MDE). Overall deformation refers to the concatenation of deformations from both the cross-modality and the intra-modality registration components. MDE is then the average length of the difference between a predicted deformation and the corresponding label deformation. It is noteworthy that for areas with roughly the same pixel values it is not even possible to predict the correct deformation and deformation on those areas is determined by the regularization used.

Detailed description of the metrics is given in the suplementary materials.

Results for the synthetic data set experiment can be seen in Table 2.

The models using the deformation equivariance encouraging losses systematically outperformed the models not using them. Four out of eight trainings with the registration component but without the deformation equivariance encouraging losses did not converge at all to a meaningful optimum. An example prediction of a such training is shown in Figure 5. The performance of the models without the deformation equivariance losses varies a lot and in few cases the performance is even quite good. However, when using either of the deformation equivariance losses the trainings always converged robustly to a meaningful optimum.

Models using both the equivariance similarity loss and the commutation loss performed the best in terms of the similarity metrics. However, the models having only either the equivariance similarity loss or the commutation loss also performed very well and it is questionable whether the differences in performance when using this kind of synthetic data set will be relevant in real world applications.

Registration accuracy is well below one pixel for all the models with deformation equivariance losses which is an excellent performance.

Results for cross-modality brain MRI synthesis experiment can be seen in Table 3.

All the three proposed variants perform very well compared to the oracle model trained with aligned data. Perhaps surprisingly on PSNR metric the model DefSim + Com + EqAdv even outperforms the oracle. The result is similar to the one obtained by Kong et al. (2021) where also the model involving registration trained with non-aligned data outperformed the oracle pix2pix model trained with aligned data. They theorize it might be due to the fact that even their aligned data set is not really aligned. However, those misalignments should then be present also in the evaluation data so it is questionable whether that explains the improvement. In our case the most likely explanation is that the predictions by the oracle model seem to be slightly sharper based on visual inspection. The loss of sharpness might actually benefit the PSNR metric. Note that all the differences are small, both metric-wise and visually. Example predictions can be seen in Figure 7.

The two configurations DefSim + DefUncondAdv, and DefSim + DefUncondAdv + Aug using the losses similar to Kong et al. (2021) perform significantly worse than our proposed models. Without augmentation the desired distributions of $F\left(x^{\left(i\right)}\right)$ and ${\tilde{y}}^{\left(i\right)}$ are very different and, as assumed, the training did not converge to anything meaningful. Together with augmentation the distributions overlap significantly and the results are clearly better, although still nowhere close to acceptable. With more difficult data, e.g, if there is a systematic posture difference between inputs and labels, designing augmentations which would work even this well could be very difficult.

Deformation accuracy is very good for all the three proposed models. Average displacement in the training data is over $10$ voxels with maximum displacement being over $32$ voxels which is a lot considering the patches with side length of $64$ voxels. On the other hand, the small patch size has the benefit that the elastic component within each patch is quite small.

Results for virtual histopathology staining experiment can be seen in Table 4.

Virtual histopathology staining data set was particularly difficult one from the viewpoint of image synthesis. The data is not very uniform and the shapes and locations of cell nuclei are in some cases impossible to predict perfectly as the input images simply do not contain information about them. However, the deformations between the inputs and the labels are quite simple.

The model with both of the deformation equivariance encouraging losses performed the best on all the metrics. However, the model DefSim + DefUncondAdv + Aug using the losses similar to Kong et al. (2021) did not perform much worse, and even outperformed the two other models proposed by us on the PSNR metric. With this data set the distributions of the desired predictions $F\left(x^{\left(i\right)}\right)$ and the training labels ${\tilde{y}}^{\left(i\right)}$ are very close to each other, i.e. there are no systematic geometrical differences. In such settings the method by Kong et al. can be expected to perform relatively well. Note also that the metrics might miss subtle improvements in geometric accuracy since no perfectly registered ground truth labels are available, and the fact that our proposed methods perform systematically better in terms of the input image based NMI metric actually suggests that to be the case. The non-augmented training DefSim + DefUncondAdv looks as good superficially as the augmented one but the predictions are clearly shifted, resulting in significantly worse metrics. Without augmentation, the distribution of translations between the inputs and the labels is not symmetric around the origin, resulting in a worse outcome.

Visual inspection reveals an interesting trend about the effects of the different equivariance encouraging losses. From our three methods the metric-wise worst performing model EqSim + EqAdv seems to be more inclided to guess nucleai whereas the two other models with the commutation loss seem to place them in more certain locations. As a result the model with only the equivariance similarity loss looks most realistic to the eye as it’s nuclei density is closer the ground truth. The effect is visible in Figure 8. We suspect this is due to the commutation loss more directly enforcing deformation equivariance which will require the shape of the nuclei to be known. The equivariance similarity loss does not have the same effect allowing the adversarial training to dominate more. This is also hinted by the fact that during the training model EqSim + EqAdv had significantly lower adversarial loss for the generator than the other two models.