SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yuanhao Wang Ramzi Idoughi Wolfgang Heidrich
King Abdullah University of Science and Technology (KAUST) {yuanhao.wang,ramzi.idoughi,wolfgang.heidrich}@kaust.edu.sa
We present a novel learning-based framework for reconstructing 3D structures from tilt-series cryo-Electron Tomography (cryo-ET) data. Cryo-ET is a powerful imaging technique that can achieve near-atomic resolutions. Still, it suffers from challenges such as missing-wedge acquisition, large data size, and high noise levels. Our framework addresses these challenges by using an adaptive tensorial-based representation for the 3D density field of the scanned sample. First, we optimize a quadtree structure to partition the volume of interest. Then, we learn a vector-matrix factorization of the tensor representing the density field in each node. Moreover, we use a loss function that combines a differentiable tomographic formation model with three regularization terms: total variation, boundary consistency constraint, and an isotropic Fourier prior. Our framework allows us to query the density at any location using the learned representation and obtain a high-quality 3D tomogram. We demonstrate the superiority of our framework over existing methods using synthetic and real data. Thus, our framework boosts the quality of the reconstruction while reducing the computation time and the memory footprint. The code is available at https://github.com/yuanhaowang1213/adaptivetensordf.
Tilt-series Cryo-ET is an important imaging tool used in a variety of fields, such as structural biology [5] and material science [48]. It consists of reconstructing a 3D tomogram from captured 2D projections of the scanned sample. In cryo-ET, due to hardware constraints, the scanning process is performed from a limited angular range, typically less than  $120^{\circ}$ . The tomographic reconstruction with missing-wedges is known to be a challenging task. In addition, sample motion and deformation that occurs during the acquisition step, induce misalignment between the projections, which deteriorates the reconstruction quality. Also, cryo-ET requires a large amount of data, which makes the reconstruction a resource-intensive task, both in terms of computational time and memory. Last but not least, to prevent the sample from getting damaged, the scanning is performed with low-intensity electron beams. This results in high level noise in the captured projections.
In the literature, several approaches have been proposed to overcome some of these challenges. The misalignment and the motion can be corrected using marker tracking [28, 35, 49], or even by tracking specific features between the projections [10, 19]. On the other hand, the high level noise was initially reduced by applying image-based or volume-based denoising approaches on the projections or the tomogram respectively. Recently, Bepler et al. [7] proposed to reconstruct two copies of the tomogram using odd/even projections only. Then, using Noise2Noise [29] technique they learn a denoised version of the tomogram. Kniesel et al. [26] combines a noise model for Scanning Transmission Electron Microscopy (STEM) and an implicit 3D shape representation in a differentiable framework to denoise the reconstructed data.
In this work, we propose a new framework for cryo-ET that jointly reconstructs and denoises the tomograms. This framework is based on a quadtree structure that we use to define an adaptive Tensorial Density Fields (TensorDF) representation. Each node of the quadtree is featured with a tensorial density fields to represent the tomogram's density at its associated region. The loss function that guides the learning process is composed of four terms: a differentiable tomographic image formation model, a total variation term that encourages smoothness in the reconstruction, a boundary consistency constraint that enforces agreement between the adjacent nodes, and an isotropic Fourier prior that penalizes directional artifacts and helps in denoising the reconstructed tomogram.
The evaluation of our approach on synthetic and real datasets shows considerable improvements comparing to state-of-the-art reconstruction methods, including some recent approaches based on neural representations. The main contributions of our work are summarized as follows:
- We propose an adaptive tensorial density fields representation, based on a quadtree structure, well-suited to large datasets. This allows us to achieve a breakthrough in 4K-resolution tomography reconstruction in less than a day.
- We propose a novel isotropic Fourier prior to remove the directional artifacts, and reduce the noise level.
Computed tomography (CT). Cryo-ET is one of the various types of tomographic inverse problem, which reconstructs a density volume of the scanned object from captured projections from different views. Analytical reconstruction algorithms, such as Filtered Back-Projection (FBP) [13] and Weighted Filtered Back-Projection (WFBP) [43] are commonly used to solve this problem. These approaches produce fast and accurate reconstructions but require a considerable number of projections that are uniformly sampled in the angular space. Algebraic Reconstruction Technique (ART) [18] and its variants, like the Simultaneous Algebraic Reconstruction Technique (SART) [2], solve the tomography problem in an iterative way. These methods are well-suited to challenging CT reconstructions (sparse view, missing wedges, noisy projections, etc.), since they can be combined with different regularizers like the total variation into optimization frameworks [21, 39, 23]. However, these approaches suffer from high computational requirements and fastidious hyper-parameter tuning. Learning-based approaches have been introduced in CT to improve reconstruction quality through pre-processing projections [3, 17, 52], post-processing the reconstructed tomogram [40, 32], or using neural networks to learn a differentiable reconstruction operator [1, 25, 20]. More recently, combining Deep Image Prior [4, 6] and implicit representations [53, 46, 45] with traditional reconstruction approaches has produced unprecedented reconstruction results. Nevertheless, these approaches are not well-designed to overcome the several challenges of cryo-ET, mentioned in Section 1, especially the high level noise.
Cryo-ET data Denoising. Cryo-ET acquisition uses low-dose beams to avoid damaging the samples, but this results in capturing noisy projections [42]. In the literature, several models have been proposed for the cryo-ET noise, such as the additive white Gaussian noise [7, 56], or a Poisson-Gaussian noise [55]. Most existing denoising approaches are applied either before or after the reconstruction step [14], which is performed using classical algorithms like WFBP. The pre-reconstruction approaches aim to denoise the 2D projections using denoising algorithms such as the bilateral filter [24], the non-local means filtering [51], the wavelet shrinkage filter [22], and deep learning-based techniques [9]. On the other hand, post-reconstruction denoising is applied directly to the tomograms to save the linearity between captured projections and reconstructed volume. This approach is unlikely to introduce new artifacts or suppress existing features. Total variation [54], and BM4D [33] were the state-of-the-art in 3D denoising, for a decade. However, these approaches require considerable computational resources. In structural biology, subtomogram averaging [8] is a specific denoising technique for data containing several copies of the same molecule. Recently, learning-based methods gained great success in denoising tasks. Specifically, unsupervised methods like Noise2Noise [29], Noise2Void [27], present a great potential for tomograms denoising, given the absence of ground truth data in cryo-ET field. Thus, the Topaz algorithm [7] produces clean tomograms by leveraging the Noise2Noise concept to train on pairs of noisy tomograms reconstructed from odd/even projections. Kniesel et al. [26] proposed to jointly learn a model for 2D sensor noise and a 3D implicit representation of the scanned sample. However, this approach depends highly on
the noise level. In our approach, the denoising task is performed during the reconstruction process through the proposed loss function and the parametrization of our adaptive representation.
3D Neural representation. Neural representation has seen impressive and rapid development in the last few years [47]. The Neural Radiance Fields (NeRF) [36] use Multi-Layer Perceptron (MLP) networks to learn a mapping between spatial coordinates and physical scene properties (e.g., density field, color). However, NeRF-like approaches, also known as neural representations, suffer from long training times and slow rendering. Techniques such as octree structures [31, 15, 45, 50], multi-scale network architecture [34], network factorization [44], caching [16], multi-resolution hash encoding [37], and tensorial-based representation [11] have been applied to handle this issue. Some works adapted these neural representations to solve CT reconstruction with impressive results [46, 53, 45, 26]. In this work, we leverage a quadtree structure to build an adaptive tensorial-based representation. We also incorporate several regularizers in our loss function to handle the high level of noise in cryo-ET data. In addition, our method can deal with large datasets (4K resolution) in an unprecedented computational time.
In this section, we introduce our approach, Adaptive TensorDF, which leverages a quadtree structure to create a multi-scale and effective tensor-based representation for reconstructing noisy tilt-series cryo-ET tomograms. Our framework represents the scanned sample's density field as a 3D continuous neural field. Since the cryo-ET tomograms have a low extension along the  $z$ -axis, we use an adaptive quadtree structure to partition the volume instead of an octree. Inside each node of the quadtree, the density field is represented using a tensorial representation. Our framework is composed of three steps: quadtree update, optimization of the tensorial representation, and 3D volume reconstruction. These steps are depicted in Figure 1. The quadtree update step involves optimizing the quadtree structure to achieve the best partition of the reconstructed volume. This is done by encouraging uniform nodes while limiting their total number. In the second step, we optimize a loss function to build a tensor-based representation of the density field inside each node. The proposed loss function incorporates a data-fidelity term and three regularizer terms, which we will detail in the following. The reconstruction step involves uniform sampling of the volume in the region of interest (ROI) and querying the density values at the sampled positions.
Figure 1: Our framework is composed of two steps: During the training step (a) we first update the quadtree structure using downsampled projections. Then, we freeze that structure and update the tensor representation in each node using the original noisy projections. The second step is the volume querying (b), where we uniformly sample the ROI, and use the learned representation to estimate the densities at the selected positions.
3.2 Image formation model
Cryo-electron tomography is a classical tomographic reconstruction problem, where the projection image captured by the sensor corresponds to the integration of density along the rays between the source and the detectors in the log space. After discretization, the image formation model for a given ray, can be written as as follows:
$$
\mathbf {b} _ {i} = \mathbf {A} _ {i} \mathbf {x} + \mathbf {n} _ {i} \tag {1}
$$
where  $\mathbf{A}_i$  represents the Radon transform operation along the ray  $i$ .  $\mathbf{b}_i$  and  $\mathbf{n}_i$  correspond to the intensity measured by the detector  $i$  and its associated noise, respectively.  $\mathbf{x}$  is the 3D density vector that we would like to reconstruct. As discussed in Section 1, the noise in real Cryo-ET acquisitions is complex and difficult to model, especially after applying the different pre-processing steps such as the intensity correction and the motion correction. For the sake of simplicity, we assume in this work a Gaussian white noise in the pre-processed projections. By regrouping all the captured rays, we deduce the data-fidelity loss from the Equation 1:
$$
\mathcal {L} _ {\text {d a t a}} (\mathbf {x}) = \frac {1}{2} \| \mathbf {K} (\mathbf {A x} - \mathbf {b}) \| _ {2} ^ {2} \tag {2}
$$
where  $\mathbf{K}$  is a binary mask to disable rays intersecting with the fiducial markers. This mask is introduced to suppress the projection artifacts caused by those markers.
3.3 Coordinate-based representation (CBR)
In tomography applications, the coordinate-based networks have been proposed to map the 3D spatial coordinates inside the volume of interest to the density field. This mapping is given by:
$$
f _ {\phi}: \mathbf {p} _ {i} \rightarrow \mathbf {x} _ {i} \quad \text {w i t h} \quad \mathbf {p} _ {i} \in \mathbb {R} ^ {3}, \mathbf {x} _ {i} \in \mathbb {R} \tag {3}
$$
where  $\mathbf{p}_i$  and  $\mathbf{x}_i$  are respectively the 3D coordinate in the volume to be reconstructed and the corresponding density.  $f_{\phi}$  corresponds to the mapping function that should be optimized. With the CBR, the captured projections are estimated by sampling positions along rays, applying the mapping to each sample, and then summing their contribution using the Radon coefficients. In common implicit representation [36, 26],  $f_{\phi}$  is chosen to be a fully connected MLP, as illustrated in Figure 2-(a). However, this representation is unsuitable for large volume sizes, as encountered in cryo-ET. To address this limitation, works like KiloNeRF [44], ACORN [34], and NeAT [45] suggested partitioning the volume of interest using uniform blocks or multi-scale octree-based structures, and assigning a smaller MLP or decoder to each partition (block or octree node) for a local density field representation. In our approach, shown in Figure 2-(b), we partition the ROI using a quadtree structure. Indeed, cryo-ET data have limited extension in the z-axis. Therefore, we base our representation on an adaptive quadtree structure of the x-y plan that we extend in the z-axis. For each quadtree node, we use a tensorial-based representation. Thus, for a given 3D point in the region of interest  $\mathbf{p}_i$ , we represent the density field using a sum of vector-matrix outer products, as follows:
$$
f _ {\phi} (\mathbf {p} _ {i}): \mathcal {D} \left(\sum_ {r = 1} ^ {R} \mathcal {V} _ {r} ^ {X} (\mathbf {p} _ {i}) \mathcal {M} _ {r} ^ {Y, Z} (\mathbf {p} _ {i}) + \mathcal {V} _ {r} ^ {Y} (\mathbf {p} _ {i}) \mathcal {M} _ {r} ^ {X, Z} (\mathbf {p} _ {i}) + \mathcal {V} _ {r} ^ {Z} (\mathbf {p} _ {i}) \mathcal {M} _ {r} ^ {X, Y} (\mathbf {p} _ {i})\right)\rightarrow \mathbf {x} _ {i} \tag {4}
$$
where  $\mathcal{D}$  is a one-layer decoder network, that converts encoded features to the output density.  $\mathcal{V}_r^X$ ,  $\mathcal{V}_r^Y$ , and  $\mathcal{V}_r^Z$  correspond to factorized vectors of the three modes  $(X,Y,$  and  $Z)$  for the  $r^{th}$  component of the tensor decomposition. Similarly,  $\mathcal{M}_r^{Y,Z}$ ,  $\mathcal{M}_r^{X,Z}$ , and  $\mathcal{M}_r^{X,Y}$  are matrix factors for the two denoted modes:  $(Y,Z)$ ,  $(X,Z)$ , and  $(X,Y)$ , respectively.  $R$  is the rank of the representation, which is a hyperparameter to be tuned according to the complexity of the reconstructed sample. Indeed, selecting a small rank can help denoise the tomogram, but it may lose some detailed features of the scanned sample.
Figure 2: In (a), we show the application of neural radiance field to tomography reconstruction, while in (b), we illustrate the general idea of our approach. Where we replaced the MLP by an adaptive differentiable tensor fields for faster optimization and better feature recovery.
The main advantage of the tensor representation is the huge reduction of the number of parameters needed to represent the volume. Consequently, this representation is well-suited to reconstruct the large-size volume that we commonly encounter in cryo-ET datasets.
Furthermore, it is important to take into account the contribution of the regions outside the ROI in the projection process, to ensure a more accurate output. However, there is no need to have a high-quality representation in such regions. We propose to use the same representation detailed in Eq. (4), but with smaller size tensors.
3.4 Regularizations
Total Variation. The Total Variation (TV) loss is commonly used in traditional tomographic reconstruction as a spatial regularizer to improve the reconstruction quality. The use of TV in neural fields representation approaches is not always straightforward, and could increase the computational complexity of the reconstruction. [45] proposed to compute the TV loss in the features space, before applying the decoder to get the density. [11] used this loss on the vector and matrix factors, to handle noise and outliers issues in regions with fewer observations. In our implementation, we also apply the TV prior on the vector and matrix factors of enabled quadtree node in the ROI. Then we average all the contributions of those nodes to compute the TV loss, which could be expressed as follows:
$$
\mathcal {L} _ {T V} = \frac {1}{N _ {E N}} \sum \alpha \cdot \operatorname {m e a n} \left(\| \nabla \mathcal {V} (\mathbf {p}) \|\right) + \operatorname {m e a n} \left(\| \nabla \mathcal {M} (\mathbf {p}) \|\right) \tag {5}
$$
where  $\alpha$  is a scaling factor. We find empirically that  $\alpha = 0.1$  produced the best results.  $\nabla \mathcal{V}(\mathbf{p})$  and  $\nabla \mathcal{M}(\mathbf{p})$  represent the gradient of the vector and matrix factors respectively.  $N_{EN}$  is the number of enabled nodes of the quadtree.
Boundary Consistency Constraint. In our proposed framework, each node of the quadtree has its own feature representation and is optimized separately. This will inevitably introduce discontinuity artifacts in the final reconstruction. To address this issue, we propose a Boundary Consistency Constraint (BCC) that penalizes the discrepancy between the features obtained from the tensor representation of neighboring nodes for sampling points on their shared edge. Our proposed BCC loss is given by:
$$
\mathcal {L} _ {B C C} = \sum_ {(n, m) \in O _ {b}} \text {m e a n} \left(\sum_ {\mathbf {p} \in \cap_ {n, m}} \| f _ {\phi} (\mathbf {p}) _ {m} - f _ {\phi} (\mathbf {p}) _ {n} \|\right) \tag {6}
$$
where  $O_{b}$  refers to all pairs of neighboring quadtree nodes,  $\cap_{n,m}$  is the set of sampling points on the boundary surface between nodes  $n$  and  $m$ ,  $f_{\phi}(\mathbf{p})_m$  and  $f_{\phi}(\mathbf{p})_n$  are the densities evaluated at the point  $\mathbf{p}$  using the tensor representation of the nodes  $m$  and  $n$ , respectively.
Isotropic Fourier Prior. By combining the quadtree structure with the tensor density fields representation, we reduced the number of parameters needed to represent the 3D volume, as well as the reconstruction time. However, a downside of the low rank tensor representation is that it favors reconstructions with structured artifacts in the form of axis aligned streaks.
In the Fourier space, these artifacts manifest themselves as high peaks along the vertical and horizontal directions. To solve this issue, we introduced in our loss an Isotropic Fourier Prior (IFP), to penalize such peaks in Fourier space. This constraint aims to limit the difference between the horizontal and vertical frequencies and the other frequencies in the Fourier domain. To do this, we first calculate the mean amplitude for each ring of the Fourier domain, which represents a given spatial frequency for all possible directions. Then, we apply a penalty to the horizontal and vertical frequencies that are much higher than this mean amplitude. Furthermore, we apply a weighting coefficient to penalize more the high frequencies that correspond mainly to the noise. In our implementation, we sample the ROI coarsely, and we query the volume at those samples. Then, we compute the Fourier transform slice by slice. Our proposed IFP loss could be expressed as follows:
$$
\mathcal {L} _ {I F P} = \sum_ {s, \zeta} w (\zeta) \left(| F _ {\phi} (0, \zeta) - \operatorname {m e a n} \left(F _ {\phi} (u, v) _ {| \left(u ^ {2} + v ^ {2} = \zeta^ {2}\right)}\right) | + \right. \tag {7}
$$
$$
\left. \left| F _ {\phi} (\zeta , 0) - \operatorname {m e a n} \left(F _ {\phi} (u, v) _ {| (u ^ {2} + v ^ {2} = \zeta^ {2})}\right) \right|\right)
$$
where  $F_{\phi}$  is the Fourier transform of  $f_{\phi}$  computed on the slice  $s.u,v$  , and  $\zeta$  are spatial frequencies.  $\mathbf{w}(\zeta)$  is the weighting coefficient given by:
$$
\mathrm {w} (\zeta) = \frac {1}{\exp \left(\zeta_ {0} - \zeta\right) + 1} \tag {8}
$$
$\zeta_0$  is a frequency parameter, to control the weight of the loss according to the frequency. Empirically, we found that  $\zeta_0 = 0.2$  provides an optimal weighting for the isotropic Fourier prior.
3.5 Adaptive tensor density field optimization
We have defined our tensor density representation and the different losses used for our optimization. By combining all those terms, our loss is given by:
$$
\mathcal {L} = \mathcal {L} _ {\text {d a t a}} + \lambda_ {T V} \cdot \mathcal {L} _ {T V} + \lambda_ {B C C} \cdot \mathcal {L} _ {B C C} + \lambda_ {I F P} \cdot \mathcal {L} _ {I F P} \tag {9}
$$
where  $\lambda_{TV}$ ,  $\lambda_{BCC}$ , and  $\lambda_{IFP}$  are the weights of the three regularizers defined previously. In the following, we describe the key elements of our framework's training step: updating the quadtree and sampling along the rays.
Figure 3: Illustration of (a) the quadtree updating process, and (b) the ray sampling strategy.
Quadrtree update. As mentioned previously, we build our tensor representation using a quadrtree structure, illustrated in Figure 3. First, we initialize our quadrtree structure from the ROI, and we disable the nodes outside this region. Then, we sample uniformly the volume of each node, and compute the standard deviation (STD) of the density values of the sampled
positions, to define the quadtree update loss. A high STD in a given node indicates that it contains fine details. Therefore, this node is more likely to be split into four child nodes. On the other hand, a node with low STD probably contains less features and represents uniform region of the scanned sample. Such node has higher chance to be merged or kept the same. We use a mixed-integer program (MIP) [41] to optimize this process, where, in each iteration, the nodes are either merged, split, or kept the same according to the update loss. During this optimization, the total number of nodes should stay lower than a fixed limit.
Furthermore, we use a coarse-to-fine strategy to speed up the quadtree update and deal with the high noise level. We first update the quadtree using down-sampled projections. During this step, we also update the tensor representation inside each node. This strategy reduces the impact of noise on the quadtree update. After some epochs, we keep the quadtree structure fixed and continue optimizing the tensor representation of each node using the original projections.
Ray sampling. During the optimization, we sample each ray to get a list of 3D positions to be used for the density integration and the loss evaluation. The applied sampling is not uniform, but follows a stratified random sampling strategy that takes into account on the quadtree structure. For each node that the ray goes through, we pick  $\mathbf{N}_q$  3D positions on the ray, randomly sampled from  $\mathbf{N}_q$  uniform segments of the ray. We define  $\mathbf{N}_q$  as follows:
$$
\mathrm {N} _ {q} = \left\lceil \mathrm {N} _ {\max } \frac {l _ {q}}{d _ {q}} \right\rceil \tag {10}
$$
where  $\mathrm{N}_{\mathrm{max}}$  is a hyperparameter corresponding to the maximum number of samples per node,  $l_{q}$  is the length of the ray inside the quadtree node  $q$ , and  $d_{q}$  is the length of the diagonal of  $q$ .