SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hao Zheng Hongming Li Yong Fan
hzheng1@upenn.edu, {hongming.li,yong.fan}@pennmedicine.upenn.edu
Accurate reconstruction of cortical surfaces from brain magnetic resonance images (MRIs) remains a challenging task due to the notorious partial volume effect in brain MRIs and the cerebral cortex's thin and highly folded patterns. Although many promising deep learning-based cortical surface reconstruction methods have been developed, they typically fail to model the interdependence between inner (white matter) and outer (pial) cortical surfaces, which can help generate cortical surfaces with spherical topology. To robustly reconstruct the cortical surfaces with topological correctness, we develop a new deep learning framework to jointly reconstruct the inner, outer, and their in-between (midthickness) surfaces and estimate cortical thickness directly from 3D MRIs. Our method first estimates the midthickness surface and then learns three diffeomorphic flows jointly to optimize the midthickness surface and deform it inward and outward to the inner and outer cortical surfaces respectively, regularized by topological correctness. Our method also outputs a cortex thickness value for each surface vertex, estimated from its diffeomorphic deformation trajectory. Our method has been evaluated on two large-scale neuroimaging datasets, including ADNI and OASIS, achieving state-of-the-art cortical surface reconstruction performance in terms of accuracy, surface regularity, and computation efficiency.
The analysis of the cerebral cortex using magnetic resonance imaging (MRI) is crucial in understanding neurodegenerative diseases [9, 43] and psychological disorders [42]. Since the cerebral cortex is a thin (a few millimeters thick) and highly folded sheet between the inner (white matter: WM) and outer (pial) surfaces, voxel-based segmentation methods cannot accurately capture its complicated morphology [17]. Instead, triangular meshes have been widely used for cortical surface reconstruction (CSR) [14] in order to accurately measure cortical thickness, volume, and gyration. Although well-established methods for CSR (e.g., FreeSurfer [17], BrainSuite [48]) can produce promising results, they often require significant computational resources (e.g., 6h/subject [17]) and may necessitate manual editing in order to attain sub-voxel precision.
Recently, deep learning (DL) methods have achieved significant improvement in CSR in terms of both accuracy (sub-voxel error) and efficiency (orders of magnitude faster) [10, 13, 19, 20, 26, 30, 31, 41, 47, 59]. These DL methods can be broadly classified into two categories according to the representation of the output surfaces. (I) Implicit surface representation, such as signed distance function [13, 19], occupancy filed [13], and level set [41], can be predicted by neural networks, and cortical surface is then generated by a marching cubes algorithm [27]. (II) Explicit surface reconstruction methods take a coarse or fine initial mesh as input and directly predict a target mesh [10, 20, 26, 30, 31, 59]. A detailed comparison of the existing DL-based CSR methods is summarized in Table 1.
Table 1: A comparison of existing DL-based CSR methods and ours. For input, "volume": the whole 3D volume; "cube": a sub-volume; FS: fine surface mesh; CS: coarse surface template; IS: implicit surface representation. For output, explicit surface: triangular mesh; implicit surface: signed distance function or occupancy field.
Methods	Input
Img.	Surf.	Network	Output
Explicit surf.	Joint surf. reconstruction	Direct thickness estimation
CNN	GNN/MLP
DeepCSR [13]	volume	✓	✓
vox2surf [19]	vol.+cube	✓	✓
FastCSR [41]	volume	✓
PialNN[31]	cube	FS	✓	✓	✓
CorticalFlow [26, 47]	volume	CS	✓	✓
vox2cortex [10]	volume	CS	✓	✓	✓	✓
TopoFit [20]	volume	CS	✓	✓	✓
CortexODE[30]	cube	FS	✓	✓	✓
Ours	volume	FS+IS	✓	✓	✓	✓
However, the existing DL-based CSR methods are subject to limitations. First, the interdependence between the inner and outer cortical surfaces is generally ignored, and therefore separate or multi-stage DL models are typically trained to reconstruct both the inner and outer cortical surfaces and intersections between them may occur. Even if both the inner and outer surfaces can be reconstructed simultaneously [10], they are loosely combined with no topological constraints. Second, complex DL architectures are commonly used, with separate learning of image and surface/vertex features using both convolutional neural networks (CNNs) and graph neural networks (GNNs)/multi-layer perceptrons (MLPs). Moreover, graph convolutions become less scalable as the number of vertices in the mesh grows to accommodate complex shapes, and often fail to learn diffeomorphic mappings to produce genus-zero regular meshes [10]. Third, using a coarse mesh template incurs difficulty of learning large deformations for highly folded cortical regions and may lead to non-smooth deformation and undesirable artifacts. Last but not least, the cortical thickness estimation is neglected in all existing DL-based CSR methods (i.e., needs a separate step to compute cortical attributes as in conventional pipelines [17, 56]). It may serve as an anatomical constraint to couple the inner and outer cortical surfaces, help generate cortical surfaces with topological correctness, and facilitate quantitative analysis of the cortical thickness.
In order to robustly reconstruct the cortical surfaces with topological correctness, we develop a DL-based approach to simultaneously reconstruct both the inner and outer cortical surfaces and estimate the cortical thickness by optimizing and deforming an initialized midthickness surface. First, our method explicitly couples the inner and outer surfaces by jointly learning three diffeomorphic flows to optimize the initialization midthickness surface to lie halfway between the inner and outer cortical surfaces and deform it to inner and outer cortical surfaces, respectively. Second, instead of designing a complex mixed architecture of CNNs and GNNs/MLPs, our method employs a single model of 3D CNNs to predict the diffeomorphic flows from a multi-channel input, consisting of a 3D brain MRI, a ribbon segmentation map that encodes structural information of the cerebral cortex, and a signed distance function that implicitly encodes the initialization surface. Third, our method calculates the diffeomorphic deformation trajectories in a continues coordinate space rather than on a 3D voxelwise grid, achieving higher sub-voxel accuracy from fine-grained velocity fields while maintaining reasonable computational efficiency even as the number of mesh vertices increases. Fourth, we devise an efficient and reliable approach to initialize a fine midthickness surface from a cortical ribbon segmentation result, followed by topology correction to ensure genus zero. Finally, a vertex-wise thickness estimation can be obtained by tracing the geodesic trajectory of each vertex during the mesh deformation process. In summary, our new DL framework differs from the existing DL-based CSR approaches in its coupled reconstruction of multiple surfaces and the simultaneous cortical thickness estimation, which facilitates robust reconstruction of the cortical surfaces with spherical topology. Ablation studies and comparison experiments on two large public datasets (ADNI [22] and OASIS [33]) have demonstrated that our method attains superior performance over the state-of-the-art methods [10, 13, 26, 30, 31, 47].
2.1 Learning-based Cortical Surface Reconstruction
Recent years have witnessed a surge of interest in geometric DL-based methods for general computer vision tasks [15, 29, 49, 52, 54, 57, 58]) and biomedical object reconstruction [24, 25, 37, 38, 55].
However, their applications to biomedical tasks are limited to organs with simple shapes such as the liver and heart. The cerebral cortex, on the other hand, is a highly folded, thin structure with a significantly complex shape, necessitating more advanced approaches.
Table 1 summarizes two categories of CSR methods, i.e., implicit and explicit surface reconstruction. The implicit surface reconstruction methods typically learn a function that maps 3D coordinates to a continuous implicit representation of the surface, such as signed distance function [39] and occupancy field [34]. Compared to well-established pipelines [17], these methods have significantly improved inference efficiency (reducing time cost from hours to minutes). Despite their ability to generate surfaces at any desired resolutions [13, 18, 19, 41], these methods often rely on a marching cubes algorithm [27] to obtain a triangular mesh and post-processing topology correction methods [7] for ensuring correct spherical topology of the reconstructed surfaces.
The explicit surface reconstruction methods typically learn a mapping from an initialization mesh to a target surface and the mapping can be modeled as a deformation model [52, 53, 55]. Particularly, Pixel2Mesh [52] utilizes GNNs to learn vertex-wise deformations of an ellipsoid and increase the mesh resolution in a coarse-to-fine manner. Such a strategy has also been adopted in Voxel2Mesh [55] to reconstruct simple human organ surfaces (e.g., liver, hippocampus) from CT/MR images. Similarly, PialNN [31] learns to deform an input inner cortical (i.e., WM) surface to a target outer (pial) surface by a sequence of deformation blocks; TopoFit [20] warps a topologically-correct template surface to fit the WM surface by a series of graph convolutional blocks; and Vox2Cortex [10] deforms a brain template surface to target cortical surfaces by leveraging combined CNNs and GNNs. Another line of research leverages ODE to parameterize the deformation of vertices as diffeomorphic flows, such as CoticalFlow methods [26, 47] and CortexODE [30]. However, all these methods may produce intersecting inner and outer cortical surfaces because they predict the inner and outer cortical surfaces either using separate DL models or using a joint DL model without explicit constraints to penalize the generation of intersecting surfaces.
2.2 Diffeomorphic Deformation
A diffeomorphism is theoretically differentiable and invertible and guarantees smooth and one-to-one mapping [46]. Diffeomorphic deformation has been widely used in image registration [2, 4, 5, 35, 51]. It can be generated from a velocity field  $\mathbf{v}$  by integrating an ordinary differential equation (ODE) [1],
$$
\frac {d \phi (\mathbf {x} , t)}{d t} = \mathbf {v} (\phi (\mathbf {x}, t), t), \text {a n d} \phi (\mathbf {x}, t) = \phi (\mathbf {x}, 0) + \int_ {o} ^ {t} \mathbf {v} (\phi (\mathbf {x}, t), t) d t, \tag {1}
$$
where  $\phi (\mathbf{x},0) = \mathbf{x}$ . The velocity field can be stationary [4] or time-varying [8]. Standard numerical integration techniques, such as the Euler method and the Runge-Kutta method [11], can be used to perform the integration. CoticalFlow methods [26, 47] and CortexODE [30] parameterize the ODE by a neural network [12]. CoticalFlow requires the deformation field to be Lipschitz at each time step to ensure a bijective mapping with Lipschitz inverse and derives a stability condition for the numeric approximation of  $\phi$ . As a chain of deformation modules is trained in sequential stages to predict a series of deformation fields that deform the initial mesh template to the target surface in a coarse-to-fine manner, there is a risk of generating self-intersections due to the low-resolution template in the coarsest level model and the whole process is prone to error accumulation in multiple steps. CortexODE [30] is built upon NODE [12] and operates on a topologically-corrected high-resolution initialization mesh. The sufficient condition of diffeomorphism can be satisfied if the deformation network is Lipschitz continuous and a sufficiently small step size is used in numerical approximation. Most learning-based image registration approaches [6, 28, 36] adopt stationary velocity field (SVF), and the integral (displacement field) is computed on a voxel-grid using the scaling and squaring method [3], yielding comparable registration accuracy to conventional methods while significantly improving efficiency. Our work lies at the intersection of these methods in that our method uses CNNs to parameterize multiple SVFs based on a multi-channel input and train the SVFs jointly to optimize cortical surfaces under topology-preserving and inverse-consistent transformation regularizations.
As illustrated in Fig. 1, our framework consists of two parts: a pipeline that estimates a midthickness surface, represented both explicitly as a 3D mesh and implicitly as a 3D distance map (Sect. 3.1), and
(a) Surface initialization
Figure 1: Overview. (a) Surface initialization. An FCN generates cortical ribbon segmentation maps from the raw image. An SDF is derived, followed by topology correction and marching cubes algorithms, and then used for initializing the midthickness surface. (b) Cortical surface reconstruction network. It takes a multi-channel input and learns three parallel velocity fields (VFs) in parallel.  $\mathrm{VF}_1$  is used to generate a deformation field for optimizing the midthickness surface  $S_M$ .  $\mathrm{VF}_2$  (or  $\mathrm{VF}_3$ ) is used to deform  $S_M$  to the WM surface  $S_W$  (or pial surface  $S_G$ ) and then  $\mathrm{VF}_3$  (or  $\mathrm{VF}_2$ ) deforms  $S_W$  (or  $S_G$ ) to midthickness surface  $S_M''$  (or  $S_M'$ ) to constrain topology between cortical surfaces.
(b) Cortical surface reconstruction network
an end-to-end fully convolutional network (FCN) that reconstructs multiple surfaces and estimate cortical thickness simultaneously (Sect. 3.2). Loss functions are presented in (Sect. 3.3).
3.1 Midthickness Surface Initialization
The human cerebral cortex is a 2D sheet with an average thickness of  $\sim 2.5mm$  and has a highly folded geometry with peaks (i.e., gyri) and grooves (i.e., sulci) [17]. The existing CSR methods typically use a smoothed template surface estimated from a group of subjects [10, 47], or a WM surface [30, 31] to initialize the surface reconstruction. However, our preliminary experiment shows that the closer the initial surface is to its target surface, the higher the reconstruction accuracy is (see Fig. 2(a)). Thus, we propose to extract the midthickness layer which lies halfway between the inner and outer cortical surfaces as an initialization surface, which brings three advantages: (1) Deforming a surface from the midthickness surface reduces the learning difficulty and improves accuracy by avoiding learning "large" deformations. It also strikes a balance between optimizing the inner and outer surfaces, making it less challenging. (2) The distance between the midthickness surface and the pial surface can help prevent topology errors in cases of deep sulci, as this layer still provides a clear separation between them. (3) Deforming the midthickness surface inward and outward to the inner and outer surfaces establishes a one-to-one mapping that explicitly encodes the correspondence between surfaces, facilitating coupled surface learning and improving CSR accuracy.
A key challenge is how to obtain a surface as close to the midthickness surface as possible. A straightforward method is to inflate the WM surface along the normal direction by a pre-determined distance [30]. However, this method cannot accurately estimate the midthickness surface in that the cortical thickness varies across the cortex and the normal is merely an approximation using neighboring faces. We propose a more accurate method by leveraging the cortex ribbon segmentation map (i.e., the filled interior area of WM and pial surfaces). Given an input brain MRI volume  $I \in \mathbb{R}^{D_1 \times D_2 \times D_3}$ , we utilize a 3D U-Net [44] to generate a WM segmentation map  $M_W \in \mathbb{R}^{D_1 \times D_2 \times D_3}$  and a GM segmentation map  $M_G \in \mathbb{R}^{D_1 \times D_2 \times D_3}$  (see Fig. 1(a)). The network is trained on a large-scale public neuroimaging dataset [22] by minimizing the cross-entropy loss between the prediction and ground truth which can be obtained using existing pipelines [17, 45]. Based on the predicted segmentation map  $M_W$ , we generate a signed distance function (SDF),  $K_W \in \mathbb{R}^{D_1 \times D_2 \times D_3}$ , using a distance transform algorithm:  $d(v_i) = SDF(v_i)$  is the minimal Euclidean distance of voxel  $v_i \in I$  to the boundary voxels. Voxels with values equal to zero represent the surface boundaries and voxels with negative or positive values encode their distances to the surface boundaries inward or outward, respectively. Similarly, we generate an SDF,  $K_G$ , for the pial (gray matter) surface. A new SDF can be obtained by averaging the WM and GM SDFs:  $K_M = (K_G + K_W) / 2$ , whose 0-level defines the midthickness surface implicitly. To ensure the midthickness surface maintains a spherical topology, a fast topology check and correction algorithm [7, 30] is then applied to the implicit surface  $K_M$ .
Figure 2: (a) The relationship between the initialization surface position (X-axis) and the surface reconstruction accuracy (Y-axis) using PialNN [31]. Orange and blue lines represent WM and pial surfaces reconstruction resp. (b) Illustration of a symmetric deformation trajectory and cortical thickness estimation. (c) Illustration of DDM. Each vertex is deformed by sampled velocities.
Finally, the initialization midthickness surface is extracted by the marching cubes algorithm [27] from the 0-level of  $K_{M}$  and parameterized by a triangular mesh  $S_{0}$ .
3.2 Coupled Reconstruction of Cortical Surfaces
Based on the formulation in Eq. 1, an accurate parameterization of the SVFs is crucial to modeling the diffeomorphic trajectory of each vertex and reconstructing the cortical surfaces. In this section, we will show how our method achieves the goals.
Feature Extraction from a Mutli-channel Input of Brain MRI, Cortical Ribbon, and Cortical Surface. The existing DL-based CSR studies have demonstrated that it is critical to fuse both image features extracted from the MRIs using CNNs and geometry features extracted from surface meshes using GNNs/MLPs. However, besides employing two different network architectures, these methods typically learn the image and geometry features separately before fusion, which does not adequately utilize image texture and surface geometric information. To better model the SVFs, we propose to learn features from a multi-channel input consisting of a brain MRI, its cortical ribbon segmentation maps, and its midthickness surface represented as an SDF:  $I_{comb} = I \odot M_{W \oplus G} \odot K_M$ , where  $\odot$  is channel-wise concatenation and  $M_{W \oplus G}$  is a multiclass mask (BG = 0, WM = 0.5, GM = 1).
Such a feature learning procedure brings two key benefits: (1) Utilizing heterogeneous features enables mutual knowledge distillation. The brain MRI contains detailed texture and semantic information but may include noise and irrelevant regions far from the target surfaces. The cortical ribbon segmentation maps contain structural/semantic information about the cortical sheet and can act as an attention guide for extracting informative features around its boundaries. The SDF implicitly embeds the surface location information and relative relation between all voxels. Together, the multi-channel input provides richer and complementary information for our model to reconstruct the surfaces. (2) Only a single CNN is needed, which simplifies the model design and improves efficiency. Features can be extracted in a single forward pass for all coordinates. When scaling up the number of vertices in mesh, we can interpolate in the feature space efficiently.
Coupled Learning of Cortical Surfaces. The goal is to learn diffeomorphic deformations that deform the initialization midthickness surface  $S_0 \subset \mathbb{R}^3$  to its target WM and pial surfaces,  $S_W$  and  $S_G$ . Taking into account the discrepancy between the initialization  $S_0$  and the true midthickness surfaces  $S_M$ , our method also learns a diffeomorphic deformation to optimize the initialization midthickness surface. In total, our method learns a function to model three diffeomorphic deformations  $f_\theta(I_{comb}, S_0) = (\phi_M, \phi_W, \phi_G)$ , using an FCN and several diffeomorphic deformation modules (DDMs). Specifically, the FCN has a similar architecture as U-Net [44], consisting of a 5-level hierarchical encoder-decoder with skip connections as shown in Fig. 1(b) (see Supplementary Materials for details). By taking the multi-channel input, the FCN learns to estimate three dense SVFs jointly, denoted by  $\mathbf{v}_M$ ,  $\mathbf{v}_W$ , and  $\mathbf{v}_G$ . We then use  $\mathbf{v}_M$  to compute  $\phi_M$  that deforms  $S_0$  to the true midthickness surface  $S_M$ ,  $\mathbf{v}_W$  to compute  $\phi_W$  that deforms  $S_M$  inward to the WM surface  $S_W$ , and  $\mathbf{v}_G$  to compute  $\phi_G$  that deforms  $S_M$  outward to the pial surface  $S_G$ . By doing so, we establish a one-to-one mapping across  $S_W$ ,  $S_M$ , and  $S_G$ .
However, since  $\phi_W$  and  $\phi_G$  are computed using different SVFs, they may cause non-invertible transformation around the midthickness surface. To address this issue, we utilize the property of diffeomorphic mapping to compute a symmetric deformation trajectory of each vertex and devise a
symmetric cycle function  $\mathcal{L}_{cyc}$  (Eq. 2) for training. Fig. 2(b) illustrates the deformation of a vertex  $\mathbf{p}_{Mid}$  outward to  $\mathbf{p}_{GM}$  using  $\phi_G(\mathbf{v}_G)$  followed by the deformation inward to  $\mathbf{p}_{Mid}'$  using  $\phi_W(\mathbf{v}_W)$ , which should be as close to  $\mathbf{p}_{Mid}$  as possible. Similarly,  $\mathbf{p}_{Mid}''$  should also be as close to  $\mathbf{p}_{Mid}$  as possible when deformed inward by  $\phi_W$  followed by the outward deformation by  $\phi_G$ . The symmetric cycle loss is formulated as:
$$
\mathcal {L} _ {c y c} = \frac {1}{N} \sum_ {\mathbf {p} \in S _ {M}} \| \mathbf {p} _ {\phi_ {W} \circ \phi_ {G}} - \mathbf {p} \| _ {2} ^ {2} + \| \mathbf {p} _ {\phi_ {G} \circ \phi_ {W}} - \mathbf {p} \| _ {2} ^ {2}, \tag {2}
$$
where  $\mathbf{p}_{\phi_b\circ \phi_a}$  represents deforming a vertex  $\mathbf{p}$  using velocity fields  $\mathbf{v}_a$  and  $\mathbf{v}_b$  sequentially. It ensures simultaneous alignment of  $\mathbf{p}_{\phi_W\circ \phi_G}\simeq \mathbf{p}\simeq \mathbf{p}_{\phi_G\circ \phi_W}$  for each vertex  $\mathbf{p}\in S_M$ . Fig. 3 demonstrates the effectiveness of  $\mathcal{L}_{cyc}$  in the coupled reconstruction of multiple surfaces in the network optimization. Moreover, we can trace the geodesic trajectory of each vertex during the mesh deformation process for estimating vertex-wise cortical thickness (Fig. 2(b)).
Figure 3: Example vertex deformation trajectories with and without  $\mathcal{L}_{cyc}$ .
Diffeomorphic Deformation Module (DDM). To numerically solve the ODE in Eq. 1, scaling and squaring (SS) method [3] can be applied on the image grid, followed by trilinear interpolation  $(\mathrm{Lint}(\cdot))$  for vertices in continuous coordinates. However, there are three limitations: (1) Interpolation cannot guarantee the invertibility of diffeomorphic mapping; (2) Numerical errors are amplified in the two sequential steps of SS+Lint(·); (3) Computation on grid voxels including no vertices is unnecessary. Hence, we utilize the DDM to directly compute the vertex-wise integral. As shown in Fig. 2(c), we obtain the velocity vector for a vertex with coordinate  $x$  by interpolating its neighboring velocity vectors  $(:= \vec{v}_{\mathcal{N}(\mathbf{x})})$ , i.e.,  $\vec{v}_{\mathbf{x}} = \mathrm{Lint}\big(\vec{v}_{\mathcal{N}(\mathbf{x})}\big)$ . The vertex then moves to a new coordinate  $\vec{v}_{\mathbf{x}} \cdot \frac{1}{T}$  where  $T$  is the total time steps. We can obtain overall deformation using this procedure in  $T$  steps.
3.3 Loss Functions
We design multiple loss functions to optimize the geometric precision of the reconstructed surfaces and regularize the SVFs for diffeomorphic deformation.