CSR-dMRI: Continuous Super-Resolution of Diffusion MRI with Anatomical Structure-assisted Implicit Neural Representation Learning

In implicit neural representation learning, a voxel can be represented by a neural network as a continuous function. the network $\mathcal{F}_{\theta}$ with parameters $\theta$ can be defined as:

where the input $c$ is the normalized coordinate index in the voxel spatial field, and the output $I$ is the corresponding intensity value in the voxel. The network function $\mathcal{F}$ maps coordinates to voxel intensities, effectively encoding the internal information of the entire voxel into the network parameters. The network $\mathcal{F}$ with parameters $\theta$ is also referred to as the neural representation of the voxel. Our voxel coordinates are constructed based on the size of the dMRI data and normalized along each dimension to the range of [-1,1]. This network is only applicable to the super-resolution of individual dMRI images and is not suitable for multiple dMRI images. Additional auxiliary information needs to be provided.

The overall architecture of the proposed CSR-dMRI model is illustrated in Figure 1. Due to the primary focus of dMRI data on tissue microstructure and diffusion information, which provides limited anatomical information. So, we introduce T1-weighted data to provide anatomical prior information. Additionally, by incorporating frequency-domain-based loss to preserve texture and structural information in the image, the quality of the image is enhanced. The entire network consists of a latent feature extractor and an implicit function network. The latent feature extractor is used to extract features from LR dMRI and anatomical images, while the implicit function network predicts voxel intensities at arbitrary coordinates by integrating coordinates and corresponding latent feature vectors, thereby achieving super-resolution. Specifically, for a given pair of LR dMRI and anatomical images from the dataset $D=\{{I_{lr}^{i}\in\mathbb{R}^{2\times h\times w\times d},I_{hr}^{i}\in\mathbb{R}^{s_{i}\cdot(h\times w\times d)}}\}_{i=1}^{N}$, along with their corresponding HR dMRI, where $N$ represents the total number of samples, and $s_{i}$ represents the upsampling scale factor for the $i$-th sample pair, dimension 2 indicates concatenating DWI images and anatomical images along the channel. We first utilize a latent feature extractor to transform the LR dMRI and anatomical images into feature maps $v_{lr}^{i}\in\mathbb{R}^{h\times w\times d\times 128}$, where each element $v_{c}^{i}\in\mathbb{R}^{128}$ corresponding to the feature vector of the LR dMRI at coordinate $c$. For any voxel coordinate $c$ in HR dMRI, we generate the corresponding feature map $v_{hr}^{i}\in\mathbb{R}^{s_{i}\cdot(h\times w\times d)}$ by bilinear interpolation of $v_{lr}^{i}$. Subsequently, the queried coordinate $c$ and the corresponding feature vector $v_{c}^{i}$ are concatenated as input to $\mathcal{F}_{\theta}$, and the output of the $\mathcal{F}_{\theta}$ is the predicted voxel intensity $\hat{I}^{i}_{hr}(c)$ at spatial coordinate $c$. By minimizing the difference between the predicted voxel intensity $\hat{I}^{i}_{hr}(c)$ and the real voxel intensity $I^{i}_{hr}(c)$, both the latent feature extractor and implicit function network are simultaneously optimized.

In this paper, we utilize data from the Human Connectome Project (HCP) WU-Minn-Ox Consortium public dataset¹¹1https://www.humanconnectome.org. We select 100 pre-processed diffusion and T1-weighted MRI data. Out of these, 70 were used for training, 10 for validation, and 20 for testing. The whole-brain diffusion MRI data were acquired with an isotropic resolution of 1.25mm, featuring four b-values (0, 1000, 2000, 3000). For diffusion MRI data, we only utilize the data from b1000, normalized by dividing the b1000 data by the average value of 18 b0 data. Firstly, we extract 9 patches of $40^{3}$ dimension size from each sample. Subsequently, these patches were cropped to generate HR patches of $(10\times s)^{3}$ dimension size. Finally, the HR patches were downsampled to obtain LR patches of $10^{3}$ dimension size using bicubic interpolation. The scale $s$ is randomly sampled from a uniform distribution $\mathcal{U}(2,3)$.

All networks are trained using the Pytorch framework with one NVIDIA RTX A6000 GPU (with 48 GB memory). The Adam optimizer is employed for model training with an initial learning rate of $10^{-4}$, and the batch size is set to 9. Every 200 epochs, the learning rate is multiplied by 0.5. The model is trained for a total of 1000 epochs. We compare our CSR-dMRI model with some state-of-the-art methods, including traditional interpolation method: Bicubic, deep learning-based fixed-scale super-resolution methods: DCSRN[3], ResCNN3D[4], and an arbitrary scale super-resolution method: ArSSR[23]. PSNR and SSIM are used for quantitative evaluation.

To comprehensively assess the effectiveness of our approach, we conduct experiments in three aspects: the first involves in-distribution experiments, including $2\times$ and $3\times$ experiments.The second is out-of-distribution experiments where we tested the results under the $4\times$. The third involves experiments with non-integer scale factors, testing the results under the $2.4\times$.

To evaluate the effectiveness of key components in the CSR-dMRI model, we conduct ablation experiments on different components, as shown in Table 2. INR indicates whether implicit neural representations are used, as ArSSR[23], T1w indicates the inclusion of anatomical image, $\ell_{k}$ indicates the introduction of frequency-domain-based loss. The quantitative results in Table 2 indicate that different components have positively contributed to the results. Furthermore, the structural prior information provided by anatomical images effectively enhances the image quality.