NAF: Neural Attenuation Fields for Sparse-View CBCT Reconstruction

Each pixel value of a projection image results from an X-ray passing through a cubical space and getting attenuated by the media inside. We sample $N$ points at the parts where rays intersect the cube. A stratified sampling method [13] is adopted, where we divide a ray into $N$ evenly spaced bins and uniformly sample one point at each bin. Setting $N$ greater than the desired CT size ensures that at least one sample is assigned to every grid cell that an X-ray traverses. The coordinates of sampled points are then sent to the position encoding module.

A simple MLP can theoretically approximate any function [9]. Recent studies [18, 21], however, reveal that a neural network prefers to learn low-frequency details due to “spectral bias”. To this end, a position encoder is introduced to map 3D spatial coordinates to a higher dimensional space.

A common choice is the frequency encoder proposed by Mildenhall et al. [13]. It decomposes a spatial coordinate $\mathbf{p}\in\mathbb{R}^{3}$ into $L$ sets of sinusoidal components at different frequencies. While frequency encoder eases the difficulty of training networks, it is considered quite cumbersome. In medical imaging practise [26, 28], the size of encoder output is set to 256 or greater. The following network must be wider and deeper to cope with the inflated inputs. As a result, it takes hours to train millions of network parameters, which is not acceptable for fast CT reconstruction.

Frequency-domain encoding is a dense encoder because it utilizes the entire frequency spectrum. However, dense encoding is redundant for CBCT reconstruction for two main reasons. First, a human body usually consists of several homogeneous media, such as muscles and bones. Attenuation coefficients remain approximately uniform inside one medium but vary between different media. High-frequency features are not necessary unless for points near edges. Second, natural objects favor smoothness. Many organs have simple shapes, such as spindle (muscle) or cylinder (bone). Their smooth surfaces can be easily learned with low-dimensional features.

To exploit the aforementioned characteristics of the scanned objects, we use the hash encoder [14], a learning-based sparse encoding solution. The equation of hash encoder $\mathcal{M_{H}}$ is:

Hash encoder describes a bounded space by $L$ multiresolution voxel grids. A trainable feature lookup table $\mathbf{\Theta}$ with size $T$ is assigned to each voxel grid. At each resolution level, we 1) detect neighbouring corners $\mathbf{c}$ (cubes with different colors in Fig. 1(b)) of the queried point $\mathbf{p}$, 2) look up their corresponding features $\mathbf{H}$ in a hash function fashion $h$ [23], and 3) generate a feature vector with linear interpolation $\mathcal{I}$. The output of a hash encoder is the concatenation of feature vectors at all resolution levels. More details of hash function and its symbols can be found in [14].

Compared with frequency encoder, hash encoder produces much smaller outputs ($32$ in our setting) with competitive feature quality for two reasons. On the one hand, the many-to-one property of hash function conforms to the sparsity nature of human organs. On the other hand, a trainable encoder can learn to focus on relevant details and select suitable frequency spectrum [14]. Thanks to hash encoder, the subsequent network is more compact.

We represent the bounded field with a simple MLP $\mathbf{\Phi}$, which takes the encoded spatial coordinates as inputs and outputs the attenuation coefficients $\mu$ at that position. As illustrated in Fig. 1(c), the network is composed of 4 fully-connected layers. The first three layers are 32-channel wide and have ReLU activation functions in between, while the last layer has one neuron followed by a sigmoid activation. A skip connection is included to concatenate the network input to the second layer’s activation. By contrast, Zang et al. [28] use a 6-layer 256-channel MLP to learn features from a frequency encoder. Our network is $10\times$ smaller.

According to Beer’s Law, the intensity of an X-ray traversing matter is reduced by the exponential integration of attenuation coefficients on its path. We numerically synthesize the attenuation process with:

where $I_{0}$ is the initial intensity and $\delta_{i}=\|\mathbf{p}_{i+1}-\mathbf{p}_{i}\|$ is the distance between adjacent points.

We conduct experiments on five datasets containing human organ and phantom data. Details are listed in Table 1.

We compare our approach with four baseline techniques. FDK [7] is firstly chosen as a representative of analytical methods. The second method SART [2] is a robust iterative reconstruction algorithm. ASD-POCS [20] is another iterative method with a total-variation regularizer. We implement a CBCT variant of IntraTomo [28], named IntraTomo3D, as an example of frequency-encoding deep learning methods.

Our proposed method is implemented in PyTorch [17]. We use Adam optimizer with a learning rate that starts at $1\times 10^{-3}$ and steps down to $1\times 10^{-4}$. The batch size is 2048 rays at each iteration. The sampling quantity of each ray depends on the size of CT data. For example, we sample $192$ points along each ray for the 128$\times$128$\times$128 chest CT. We use the same hyper-parameter setting for hash encoder as [14]. More details of hyper-parameters can be found in the supplementary material. All experiments are conducted on a single RTX 3090 GPU. We evaluate five methods quantitatively in terms of peak signal-to-noise ratio (PNSR) and structural similarity (SSIM) [25]. PSNR (dB) statistically assesses the artifact suppression performance, while SSIM measures the perceptual difference between two signals. Higher PNSR/SSIM values represent the accurate reconstruction and vice versa.

Our method produces quantitatively best results in both human organ and phantom datasets as listed in Table 2. Both PSNR and SSIM values are significantly higher than other methods. For example, the PSNR value of our method in the abdomen dataset is 3.07 dB higher than that of the second-best method SART.

We also provide visualization results of different methods in Fig. 2. FDK restores low-quality models with notable artifacts, as analytical methods demand large amounts of projections.

Iterative method SART suppresses noise in the sacrifice of losing certain details. The reconstruction results of ASD-POCS are heavily smeared because total-variation regularization encourages removing high-frequency details, including unwanted noise and expected tiny structures. IntraTomo3D produces clean results. However, edges between media are slightly blurred, which shows that the frequency encoder fails to teach the network to focus on edges. With the help of hash encoding, results of the proposed NAF have the most details, clearest edges and fewest artifacts. Fig. 4 indicates that NAF outperforms other methods in all slices of the reconstructed CT volume.

Figure 4 shows the performance of iterative methods and learning-based methods under different number of views. It is clear that the performance increases with the rise of input views. Our methods achieves better results than others under most circumstances.

We record the running time of iterative and learning-based methods as shown in Fig. 5. All methods use CUDA [15] to accelerate the computation process. Overall, the methods spend less time on datasets with small projections (chest, jaw and foot) and increasingly more time on big datasets (abdomen and aorta). IntraTomo3D requires more than one hour to train the network. Benefiting from the compact network design, NAF spends similar running time to iterative methods and is 3$\times$ faster than the frequency-encoding deep learning method IntraTomo3D.