Reconstruct high-resolution multi-focal plane images from a single 2D wide field image

In recent years, the super resolution of 2D natural images [23, 24, 25, 26] has been widely studied. Dong et al. [23] introduced convolutional neural network into super resolution, making deep learning play a key role in super resolution. Kim et al. proposed VDSR [27], which greatly improved the effect of super resolution due to using deeper network and a strategy of learning global residuals. EDSR [24] removed unnecessary modules of conventional residual networks, and achieved better results by expanding model architecture and stabilizing training process. To obtain more realistic HR images, Ledig et al. [25] introduced GAN into super resolution and proposed SRGAN to obtain super-resolved images with richer details. But the introduction of GAN caused more artifacts. To alleviate artifacts, ESRGAN [20] introduced residual-in-residual dense blocks and removed batch normalization, which further improved the quality of reconstructed HR images.

At the same time, super-resolution methods designed for various medical images have also emerged. Wang et al. [11] used an U-Net-like GAN model to improve the resolution of low numerical aperture(NA) fluorescence microscopy images. Ma et al. [10] proposed a multi-stage and multi-supervised model PathSRGAN to expand the application of 4x/0.1-NA pathology images. Rivenson et al. [28] extended deep learning based super-resolution to mobile-phone-based microscopy images.

For 3D medical images, Pham et al. [16] proved that LR 3D MRI images can be reconstructed into high-quality HR 3D MRI images by using deep convolution neural networks, which greatly improves the speed of obtaining HR MRI images. Zhou et al. [29] improved resolution of 3D fluorescence microscopic images by using a dual super resolution model. You et al. [14] proposed a semi-supervised 3D CT super-resolution method by introducing Wassertein constraints to GAN and node constraints to loss functions. In summary, the above 3D super-resolution methods improves resolution and imaging speed of 3D medical images.

Although 3D super resolution is helpful in expanding the application of LR medical images, the acquisition of 3D LR images still requires time-consuming axial scanning. With the progress of computational imaging, deep learning based 2D-to-3D super-resolution methods are gradually emerging. Page et al. [30] proposed LFMNet to reconstruct confocal microscopy stacks from single light filed images. Wu et al. [17] proposed Deep-Z model based on GAN, which can refocus single 2D wide-field fluorescence images to obtain 3D images.

Although both MFPINet and Deep-Z can reconstruct 3D images from a single 2D wide field image, they are very different. First, MFPINet enhances the horizontal and axial resolution at the same time. So it does not require HR 2D images as input, which greatly improves the imaging speed and expands the applications of MFPINet. Second, MFPINet is designed to output 3D images at one time instead of generating single focal plane image one-by-one like Deep-Z, which improves reconstruction efficiency.

The proposed MFPINet is a GAN framework model, which consists of two parts: generator ($G$) and discriminator ($D$). Fig. 1 shows the architecture of $G$ and $D$.

The generator consists of an encoder and a decoder, which is inspired by U-Net [19] and ResNet [21]. The encoder is composed of a preprocessing module, three residual encoding modules (REM) and a stacked convolutional layers (SCL), which are expressed as the following equation:

where $x_{3}$ stands for sRGB input image, $Conv_{k}(\cdot)$ stands for a convolutional layer with $k$ channels. $P(\cdot)$ stands for parametric rectified linear unit. $SCL(\cdot)$ stands for a stacked convolutional layers composed of three covolutional layers. $REM$ represents the residual encoding module. $SCL$ and $REM$ can be expressed as equation 2 and 3:

where $x_{k}$ stands for a feature map with k channels. $AP(\cdot)$ stands for 2D average pooling, which is used to halve the width and height of the feature map. $RB(\cdot)$ represents a residual block composed of four convolutional layers. Detailed topological connections can be found in Fig. 1.

The encoder outputs feature maps of four different scales, which are $f_{1}$, $f_{2}$, $f_{3}$, $f_{4}$ from shallow to deep. Then decoder takes all output feature maps of the encoder as input for reconstruction. Specifically, the decoder consists of three residual decoding modules (RDM) and one super-resolution reconstruction (SRR) module, which are expressed as the following equation:

where $RDM(\cdot)$ represents the residual decoder module, which is used to integrate feature information of different scales. $SRR(\cdot)$ represents the super-resolution reconstruction module, which is used to reconstruct HR MFP images from multi-scale feature maps fused by $RDM(\cdot)$. $SRR$ and $RDM$ are expressed as equation 5 and 6:

where $f_{k}$ stands for the $k^{th}$ feature map output by the encoder. $NI(\cdot)$ stands for nearest interpolation at width and height dimension of feature maps. $CAT(\cdot)$ stands for the concatenation of tensors along channel dimension. $f_{fused}$ represents feature maps fused by the three $RDM$ modules. $PS(\cdot)$ represents pixel shuffle [31] layer, which is used to double width and height dimension of feature maps. $Tanh(\cdot)$ represents tanh activation. $n$ represents the number of different focal plane 2D images in the output MFP images.

The above-mentioned encoder and decoder constitute the proposed generator. Although it borrows the U-shaped structure of U-Net, it is quite different from the original and common modified U-Net. First, the common BN module is discarded because BN is reported to cause artifacts [20]. Second, residual blocks are introduced to alleviate gradient problem caused by deep networks. Third, according to the characteristics of super-resolution tasks, only three down-sampling blocks are kept. Because super-resolution tasks are different from high-level tasks such as semantic segmentation, pixels in super-resolution tasks generally only have a greater relationship with their neighboring pixels. At the same time, because the model becomes smaller, the model can process larger images at a faster speed, thereby enabling the model to quickly generate MFP images. Fourth, the post-sampling strategy is adopted instead of interpolating the input image at the beginning, which greatly reduces the consumption of memory and computational complexity. At last, the proposed generator directly outputs 3D MFP images instead of outputting refocused images [17] one by one, which greatly improves the efficiency of acquiring MFP images.

The discriminator of the proposed MFPINet is a simple classification network, and its architecture is shown in Fig. 1. In order for the discriminator to perceive the difference of all focal plane images at the same time, the designed discriminator takes 3D MFP images as input. In addition, to reduce consumption of memory and computation, 3D MFP images are converted into 2D images, and then use 2D convolution to process them. Specifically, the original 3D MFP image represented by a four-dimensional array $H\times W\times Z\times C$ is converted into a three-dimensional array $H\times W\times(Z\times C)$.

The proposed generator is designed to learn a transformation mapping from LR 2D images to corresponding 3D HR MFP images. In the training phase, generator $G_{\theta_{g}}$ with parameter $\theta_{g}$ takes a 2D LR image $X$ as the input and output 3D HR MFP images $G_{\theta_{g}}(X)$. The discriminator $D_{\theta_{d}}$ with parameter $\theta_{d}$ takes real HR MFP images and forged MFP images by $G_{\theta_{d}}$ respectively. Discriminator $D_{\theta_{d}}$ outputs a probability value for input MFP images, which represents the model’s confidence. The loss $L_{G}$ of generator is defined as following equation:

where $N$ represents the number of samples of mini-batch in each iteration. $X^{\{n\}}$ is the $n^{th}$ 2D input image. $G_{\theta_{g}}{(X^{\{n\}})}$ is the HR MFP images obtained by the generator, and $Y^{\{n\}}$ is the real HR MFP images corresponding to $X^{\{n\}}$. $\alpha$ is a hyper-parameter used to balance absolute error loss and adversarial loss, which is set to 0.01 in the experiments.

As for the discriminator loss $L_{D}$, it is defined as the following equation:

The meaning of the parameters in equation 8 is similar to that in equation 7. In order to improve the training speed and optimization stability, Adam [32] optimization algorithm is used to optimize both generator $G_{\theta_{g}}$ and discriminator $D_{\theta_{d}}$. The learning rate is set to 0.0001 and other parameters are set to the default value. More details can be found in Experiments section.

In order to evaluate the quality of images generated by MFPINet, four common metrics are used to compare the differences between the generated MFP images $I_{g}$ and the real mechanical scanning MFP images $I_{gt}$. In order to explore the reconstruction effect of different focal planes, the MFP images are not evaluated as a whole, but as individual focal plane.Specifically, the following four metrics are used.

All experiments are completed in PyTorch 1.4 [33] and Ubuntu 18.04 equipped with two NVIDIA RTX 2080ti GPUs. For SRGAN and Deep-Z, we retrained these methods using the parameters provided in original papers on our dataset.

In this part, we show the overall effect of our model, as well as the individual effect of horizontal super resolution and axial refocusing.