Estimating Neural Orientation Distribution Fields on High Resolution Diffusion MRI Scans

TL;DR

This work proposes HashEnc, a grid-hash-encoding-based estimation of the ODF field and demonstrates its effectiveness in retaining structural and textural features and achieves a 10% enhancement in image quality while requiring 3x less computational resources than current methods.

Abstract

The Orientation Distribution Function (ODF) characterizes key brain microstructural properties and plays an important role in understanding brain structural connectivity. Recent works introduced Implicit Neural Representation (INR) based approaches to form a spatially aware continuous estimate of the ODF field and demonstrated promising results in key tasks of interest when compared to conventional discrete approaches. However, traditional INR methods face difficulties when scaling to large-scale images, such as modern ultra-high-resolution MRI scans, posing challenges in learning fine structures as well as inefficiencies in training and inference speed. In this work, we propose HashEnc, a grid-hash-encoding-based estimation of the ODF field and demonstrate its effectiveness in retaining structural and textural features. We show that HashEnc achieves a 10% enhancement in image quality while requiring 3x less computational resources than current methods. Our code can be found at https://github.com/MunzerDw/NODF-HashEnc.

Figures (8)

• Figure 1: Overview of HashEnc. 1) Given point $\boldsymbol{v}$, for each resolution grid $l$, embedding vectors of the surrounding corner points are retrieved from the lookup table by hashing their grid coordinates $t_{v,l,i}$. Then, the corner embeddings are combined into one vector via linear interpolation. The final embedding vector $t_v$ is obtained by concatenating the input coordinates $\boldsymbol{v}$ and other grid vectors $t_{v,l}$. The grid is shown in 2D instead of 3D for the sake of clarity. 2) $t_v$ is fed into a SIREN and processed by a linear layer $W$ to output the spherical harmonic coefficients $\boldsymbol{c}(\boldsymbol{v})$.
• Figure 2: GFA reconstruction images from a sagittal cerebellum slice at different training times with $M=70$ gradient directions. The scale on the top right indicates the degree of anisotropy of water diffusion. Included are GFA of the training image (1 session) and the 6 session average for reference. HashEnc fits a much more detailed ODF field after significantly less training time compared to SIREN.
• Figure 3: Qualitative reconstruction examples on a small sagittal Cerebellum section, showcasing DTI images, deconvolved ODFs, and GFA uncertainty. The scale on the bottom left indicates variability in the GFA of the ODF samples. SIREN and HashEnc are trained for 10,000 and 3,000 epochs, respectively, and compared for $M=[70, 40, 20]$ gradient directions. SIREN tends to over-smooth the ODF field, particularly at $M=70$, but is more robust with fewer gradient directions. HashEnc matches the structural and textural details of the 6-session average better and exhibits less uncertainty.
• Figure 4: DTI and GFA images of an axial slice of SIREN and HashEnc methods trained on $M=70$ gradient directions. On the right side are the DTI and GFA images of the 6 session average and 1 session. For each of the SIREN and HashEnc images we report the FSIM score to the 6 session average. We also highlight a small section indicated by the red box to demonstrate the over-smoothing effect of SIREN in comparison to the other images. HashEnc shows a better structural similarity to the 6 session average, indicated both visually and by the higher FSIM score.
• Figure 5: Cerebellum DTI and GFA images of HashEnc method with different resolutions levels $n$ and lookup table sizes $2^m$. Right are the 6 session average and 1 session images. We report the FSIM score of every image to the 6 session average on the bottom right corner. Based on the FSIM score, the network configuration of $n=14$ and $m=20$ shows the best structural similarity to the 6 session average.