[Paper notes] AIR-Net: GAN-based registration model

This article is the reading notes of the paper "Adversarial Image Registration with Application for MR and TRUS Image Fusion".

The article proposes a GAN-based supervised network model AIR-Net for multi-modal registration of MRI and transrectal ultrasound (TRUS) images, in which MRI is treated as a fixed image and transrectal ultrasound image is treated as a floating image .

Model by generator $G$ and discriminator$D$ consists of two parts, in which the generator directly estimates a conversion parameterT est T_{est}from a floating image to a fixed image$T_{est}$; Then use an image sampler to use the estimated conversion parameters $T_{est}$And ground-truth conversion parameters $T_{gt}$The floating image is registered; and the discriminator is used to determine whether the image pair passes $T_{est}$Still $T_{gt}$ For alignment.

The article regards three-dimensional images as multi-channel two-dimensional images,

The detailed structure of the generator is as follows:

One hollow convolutional layer (128 channels, dilation rate of 2) is used to expand the receptive field, two convolutional layers (128 channels, step size of 2) are used to reduce the resolution, and one contains three convolutions with residual connections The residual module of the accumulation layer (128 channels), a convolution layer (the size of the convolution kernel is $1\times 1$ , the number of channels is 8) used to reduce the number of parameters, two fully connected layers are used to get the final output, the output of the first fully connected layer is 256 dimensions, and the dimension and conversion of the output of the second fully connected layer The number of parameters is the same (if it is a three-dimensional rigid matching criterion, there are 6 parameters, and if it is a three-dimensional affine matching criterion, there are 12 parameters). If there is no special statement for the above convolution layer, the size of the convolution kernel is$3\times 3$ , followed by a ReLU activation function.

The detailed structure of the discriminator is almost the same as that of the generator. The only difference is that the output dimension of the last fully connected layer is 1, and it is followed by a Sigmoid activation function.

---

判别器的损失为：
$$\mathcal{L}(D)=-{\mathbb{E}}_{T\sim p_{gt}(T)}\left[D\left(I_f,I_m\right)\right]+{\mathbb{E}}_{T\sim p_{from}(T)}[D\left(I_f,T\left(I_m\right)\right]$$
Where $I_f $ and $I_m$ are fixed image and floating image respectively, $E_{T\sim p_{gt}\left(T\right)}\left[D\left(I_f,I_m\right)\right]$ Represents the expected loss of the discriminator of the aligned MR-TRUS image pair, and$E_{T\sim p_{from}\left(T\right)}[D(I_f,T(I_m)]$ Represents the expected loss of the discriminator for randomly aligned image pairs.

生成器的损失为：
$$\mathcal{L}(G)={\mathbb{E}}_{T\sim p_{from}(T)}\left[1-D\left(I_f,T_{est}\left(T\left(I_m\right)\right)\right)+a{\Vert T_{est}-T^{-1}\Vert }^2\right]$$
where$||T_{eset}-T^{-1}|{|}^2$ is the Euclidean distance between the estimated transformation and the randomly generated transformation.

In training, a method similar to WGAN (Wasserstein GAN) is used. In order to ensure the stability of training, after each update of the parameters of the discriminator network, the parameters of the discriminator network need to be clipped to a certain range, and the value of the clipping parameter is set to 0.1. And if the generator network is not trained once, the discriminator network will be updated twice.