DC-WCNN: A deep cascade of wavelet based convolutional neural networks for MR Image Reconstruction

Let $x\in C^{N}$ be the desired image to be reconstructed from undersampled k-space measurements $y\in C^{M}$, $M<<N$, such that $y=F_{u}x$, where $F_{u}$ is the undersampled Fourier encoding matrix. The linear inversion $x_{u}=F_{u}^{H}y$ is called zero-filled reconstruction. Reconstructing $x$ from $y$ is an ill-posed problem due to sub-Nyquist sampling. The proposed method using WCNN can be formulated as the optimization problem:

$\displaystyle\underset{x,\theta}{\operatorname{argmin}}||x-f_{wcnn}(x_{u}\lvert\theta)||_{2}^{2}+\alpha||F_{u}x-y||^{2}_{2}$

(1)

Let $x_{wcnn}=f_{wcnn}(x_{u}\lvert\theta)$, where $f_{wcnn}$ is the forward mapping of WCNN parameterized by $\theta$. Here $\alpha$ is a weight factor based on the noise of the acquired data y. The above equation enforces $x$ to be approximated by the reconstruction of WCNN in image domain without any prior information about the acquired data in k-space.

We use data fidelity (DF) unit in k-space domain after WCNN unit to ensure that the WCNN reconstruction is consistent with the acquired k-space measurements. The data fidelity operation $f_{df}$ can be expressed as,

$$\hat{x}_{rec}=\begin{cases}\hat{x}_{wcnn}(k)&\ k\notin\Omega\\ \frac{\hat{x}_{wcnn}(k)+\lambda\hat{x}_{u}(k)}{1+\lambda}&k\in\Omega\\ \end{cases}$$

(2)

Here, $\hat{x}_{wcnn}=F_{f}x_{wcnn}$, $\hat{x}_{u}=F_{f}x_{u}$, $\Omega$ is the index set of known k-space data, $F_{f}$ is the Fourier encoding matrix, and $\hat{x}_{rec}$ is the corrected k-space and $\lambda\to\infty$. The reconstructed image is obtained by inverse Fourier encoding of $\hat{x}_{rec}$, i.e. ${x}_{rec}=F_{f}^{H}\hat{x}_{rec}$. The proposed cascaded architecture, DC-WCNN, is a series of $N_{c}$ such WCNN and DF units which can be formulated as,

	$\displaystyle x_{n}$	$\displaystyle=WCNN_{n}(x_{n-1}^{df})+x_{n-1}^{df}$		(3)
	$\displaystyle x_{n}^{df}$	$\displaystyle=DF_{n}(x_{n})$		(4)

Here $WCNN_{n}$ and $DF_{n}$ denote the $n^{th}$ WCNN reconstruction block and DF unit respectively, $n=1,2..N_{c}$, $x_{0}^{df}=x_{u}$ and $x_{rec}=x_{N_{c}}^{df}$ is the output of the last DF unit.

The 2D Discrete Wavelet Transform (DWT) and inverse wavelet transform (IWT) layers of WCNN are given by,

	$\displaystyle(X_{1},X_{2},...,X_{K})=DWT(X_{in})$		(5)
	$\displaystyle X=IWT(X_{1},X_{2},...,X_{K})$		(6)

Here, $X_{in}$ be the input image or an intermediate feature map in the network, K is the number of subbands (K $=$ 4 in our case, with approximation, horizontal, vertical and diagonal subbands), $X_{k},k=1,2,...K$ denote the coefficients of the $k^{th}$ subband and X is output of wavelet recomposition.

Several studies have been done to bring wavelet transforms into CNNs ([9], [12]) . The proposed WCNN is a variant of U-Net with contraction and expansion paths wherein DWT and IWT are performed in place of pooling and unpooling layers respectively. The DWT subbands are stacked and passed to the convolutional layers. In this way, all the subbands are jointly learnt along with the inter-dependencies between them thereby enhancing spatial context. We have used residual connections wherein feature maps from contracting paths are element-wise added to the respective expanding path feature maps instead of concatenations. The residual learning strategy simplifies the optimization process and minimizes degradation of original signal information.

The WCNN as a standalone architecture unit could de-alias the US MRI image in a single step. The DC-CNN proposed by Schlemper et al. showed that a deep cascade of CNNs is similar to unfolding the optimization process of CSMRI and improves the performance of MRI reconstruction. We therefore embed WCNN as the reconstruction block in deep cascaded mode with DF units interleaved to form the DC-WCNN architecture.

The proposed DC-WCNN has cascades of WCNN and DF units (Fig.2). The WCNN consists of three WPT levels (i.e. three subsampling layers in the network). The subband images obtained from each level of DWT is fed to a 4-layer fully convolutional (FC) block in the contraction path. In the expansion path, IWT is performed after every FC block. The convolution layers have 3x3 convolution filters followed by batch normalization (BN) and rectified linear unit (ReLU) operations. In the last layer, only convolution operation is performed without BN and ReLU, to predict the residual subbands. In the DC-WCNN, the DF layers are inserted between WCNN units to correct the accumulation of possible distortions in the predicted k-space data over the cascaded blocks. We have used L2 loss function for both standalone and cascaded modes. Given the training data D consisting of a number of US and FS images as input-target pair $(x_{u},x_{t})$ as given by,

$\displaystyle L(\theta)=\sum_{(x_{u},x_{t})\in D}\|x_{t}-x_{pred}||_{2}^{2}$

(7)

Here $x_{pred}=x_{wcnn}$ in the standalone mode and $x_{pred}=x_{rec}$ in the deep cascade mode.

We have conducted experiments on the publicly available Kirby21 dataset [13] with human brain data. The dataset consists of 5460 slices of size 256x256 taken from 42 T1-weighted MPRAGE volumes out of which, 3770 slices from 29 volumes are used for training and 1690 slices from 13 volumes for validation. A fixed cartesian undersampling mask with ten lowest spatial frequencies and remaining following a zero-mean Gaussian distribution is chosen. The acceleration factor for the mask is 5x (20%) to retrospectively generate undersampled MRI images for training and testing. Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), High Frequency Error Norm (HFEN) and Normalised Mean Square Error (NMSE) metrics are used to evaluate the reconstruction quality. Wilcoxon signed-rank test with an alpha of 0.05 is used to assess statistical significance.

We have trained WCNN as standalone and within the deep cascade framework (DC-WCNN). All the standalone models are trained for 150 epochs on Nvidia GTX-1070 GPUs. In the DC-WCNN, the weights of each WCNN are initialized with those of the standalone WCNN and subsequently weights are fine tuned. The same strategy is followed for other models in deep cascade mode. The models are implemented in PyTorch and code is publicly available ¹¹1https://github.com/sriprabhar/DC-WCNN. Adam optimizer is used with a learning rate of $0.001$.

We conduct two sets of experiments, one in standalone mode and the other deep cascade mode.

Standalone mode: We compare WCNN with vanilla CNN [3] and U-Net [4] (which commonly has max pooling layers) and U-Net with mean pooling instead of max pooling layers (we refer to it as U-NetMean in Table 1). Table 1 standalone mode shows our WCNN model outperforms all other methods with best values for PSNR, SSIM, NMSE and HFEN. We note that WCNN has only 3 subsampling levels and performs better than the two U-Net architectures which have four pooling (sub-sampling) layers each. Visual comparison in Fig. 3 shows that that WPT based sub-sampling in WCNN recovers fine details better than CNN and U-NetMean (U-Net not shown as it performs same as U-NetMean). The approximation subband obtained from DWT is similar to the average pooling performed in U-NetMean, which is basically a smoothing operation. By including the detail subbands into the convolution layers, emphasis on the blurring induced by smoothing is reduced, thereby helping convolutional layers to learn kernels that maximize the overall performance of WCNN.

Deep cascade mode: In this mode, we compare DC-WCNN with DC-CNN [7] and DC-UNet [8]. From Table 1 deep cascade mode we make two observations. Firstly, the deep cascade mode boosts up the quantitative metrics as compared to standalone mode. This shows that deep cascaded architectures with data fidelity units provide better reconstruction. Secondly, our DC-WCNN model inherits the benefits of WCNN and outperforms other methods quantitatively. Visual comparison in Fig. 4 shows that DC-CNN (with two cascades of WCNN and DF) and DC-UNet produce smudged structures whereas our method recovers structures much closer to the target.

In both the modes, the metrics are found to be statistically significant (p$<$0.05). We also note that increasing the number of WPT levels increases the computational cost and hence we have chosen three levels .