DeepCSR: A 3D Deep Learning Approach for Cortical Surface Reconstruction

In order to ease the presentation, we focus on the learning of a single target surface, but this framework can be generalized to multiple surfaces by introducing mild modifications. Without loss of generality, one can implicitly define a Lipschitz surface $S\subset\mathbb{R}^{3}$ by the $\ell$-level set of a continuous function $f_{S}:\mathbb{R}^{3}\mapsto\mathbb{R}$ as,

where the function $f_{S}$ maps to a scalar for every point $p\in\mathbb{R}^{3}$ in the 3D Euclidean space. Furthermore, given $f_{S}$, the mesh representation of $S$ can be obtained by iso-surface extraction algorithms like marching cubes [20] or ray casting [44]. Therefore, reconstructing a specific surface can be thought as modeling a function such that its $\ell$-level set defines the desired geometry.

For a given surface $S$, the function $f_{S}$ can be modeled, for instance, using occupancy field or signed distance function. The former is a binary valued function defined as,

where $\mathds{1}_{pred}$ is the indicator function evaluating to one, if its predicate $pred$ is true, and zero otherwise. In this case, the level set of interest to retrieve the surface is $\ell=\frac{1}{2}$. The occupancy field $r_{S}^{occ}$ divides the 3D space in the interior $S_{int}$ and exterior of the oriented surface $S$. On the other hand, the signed distance function consists of the Euclidean distance between every point in the 3D space and its projection, $\text{proj}_{S}(p)=\mathop{\textrm{argmin}}_{q\in S}\left\lVert p-q\right\rVert_{2}$, on the surface $S$. Mathematically,

Therefore, ${r}_{S}^{sdf}$ assigns positive distances to points in the interior of the surface $S_{int}$, zero to points on the surface $S$ (i.e., $\ell=0$), and negative distances otherwise. These tools are widely used in surface modeling, then we evaluate both in the context of cortical surface reconstruction.

Our final goal is to reconstruct surfaces according to an observed MR image. Rather than conditioning the function $f_{S}$ on a specific surface, one should condition it on an input 3D MR image. While this formulation seems logical, we notice that MR images may be defined in different coordinate systems producing an unbounded input space for learning. We overcome this difficulty by co-registering the input MR images to a brain template constraining the space where the target surfaces exists to the coordinates of the bounding box containing the template brain in the template coordinate system. Therefore, our problem boils down to model the function,

where $\Omega\subset\mathbb{R}^{3}$ is a closed and bounded subspace where the template brain is defined and ${\cal I}$ denotes the space of MR images $I$ represented by a 3D grid of voxels.

Using a machine learning approach to model such a function, we first define a dataset ${\cal D}=(I_{i},S_{i})_{i=1,\ldots,N}$ of $N$ MR images $I_{i}$ and their corresponding surfaces $S_{i}$ which can be obtained using traditional pipelines for automatic cortical surface reconstruction [5, 40, 4]. Then, we parametrize the function $f_{\theta}$ in terms of learnable parameters $\theta$ and minimize the regularized empirical risk on the dataset ${\cal D}$. This learning problem can be stated as,

where ${\cal R}$ is some regularization function on the parameters $\theta$ and $\Delta$ is an appropriate cost function measuring discrepancies between the predicted implicit surface representation $f_{\theta}(p,I_{i})$ and the true implicit surface representation $r_{S_{i}}^{*}(p)$ for the given point $p\in\Omega$, MR image $I_{i}$, and its corresponding surface $S_{i}$. In the case of occupancy field, we minimize the binary cross-entropy classification loss, while for signed distance function we opt for the $L^{1}$-loss.

In order to use this approach for cortical surface reconstruction, we extend the scalar-valued functions $f_{\theta}$ to a vector-valued function $f_{\theta}\in\mathbb{R}^{4}$ where the predicted values are the surfaces representations for the inner and outer cortical surfaces further divided into left and right brain hemispheres. Consequently, the learning problem descried by Equation 5 is solved jointly for these four cortical surfaces across all of the MRIs in the dataset ${\cal D}$. This formulation is an instance of multi-task learning which often produces models with good generalization ability [36].

In practice, we approximate the integral in Equation 5 by a finite sum of sampled locations $p$ in the reference space $\Omega$. We observed that the sampling scheme employed is critical to the quality of the reconstructed surface. More specifically, as shown later in Section 4.1, the proposed framework reconstructs more accurate cortical surfaces when we sample more points near the surface, in addition to uniformly on the template space $\Omega$. In order to sample points near the target surface, we first sample faces proportionally to their area, then we sample points uniformly in these faces using the triangle point picking method [46]. Finally, we perturb these sampled points on the surface by adding a Gaussian $\mathcal{N}(0,1)$ centered at zero with $1~{}mm$ variance getting points near the surface. Such a sampling scheme provides global information to correctly align the predicted surface to the brain and a lot of local information to capture the highly folded geometry of the brain cortex.

We implement the parameterized function $f_{\theta}$ as an encoder-decoder neural network. The encoder takes as input a point $p$ represented by its coordinates in the MNI105 space and a co-registered MR image as a 3D voxel grid $I$. It first processes the input MR image by a sequence of 3D convolutional layers (Conv3D), ReLu activation functions, max pooling layers (Max3D), and fully connected layers (FC) producing multiple feature maps as indicated in the top of Figure 2. The convolutional layers use $3^{3}$ kernels, stride equal to two, and padding equal to one, with the exception of the first convolutional layer where the stride is equal to one. Our goal is to produce a rich hierarchy of visual features with different level of details. Then, inspired by Hariharan et al. [10], we project the input point coordinates into each feature map, linearly interpolate a feature value at these projected locations, and concatenate these interpolated values to form a 512-dimensional hypercolumn feature vector $c(I,p)$. This hypercolumn vector holds global and local visual cues for predicting the surface representation at the input point $p$ since it collects features from different levels of our hierarchy of features maps. As shown later in Section 4.1, this architectural design is critical to capture the high frequency details on the target surfaces.

On the other hand, the decoder receives as input the point coordinates $p$ and its corresponding hypercolumn vector $c(I,p)$ generated by the encoder for the image $I$ as previously described. The decoder processes these inputs through a sequence of fully connected layers, conditional batch normalization layers (CBN) [29] and ReLU activation functions as also indicated in Figure 2, outputting the predicted implicit surfaces representation $f_{\theta}(p,I)\in\mathbb{R}^{4}$ for the inner and outer cortical surfaces divided into left and right brain hemispheres. The decoder network uses skip-connections and follows the method proposed by Mescheder et al. [25]. However, our CBN layers compute a non-linear function conditioned on the hypercolumn vector $c(I,p)$ which encodes global and local visual cues for the prediction of the desired surfaces representation.

We train such a model from scratch by optimizing the objective in Equation 5 extended for multiple surfaces using stochastic gradient descent. More specifically, we use Adam optimizer with an initial learning rate equal to $10^{-4}$ and back-propagate the loss computed on mini-batches of 5 MR images and 1024 sampled points per image from a precomputed pool of 4 million points whereby 10% of them are sampled uniformly in $\Omega$ and 90% are sampled near the target surfaces as explained in Section 3.1.

DeepCSR receives as input a MRI scan and outputs mesh representations for the outer and inner cortical surfaces further divided into the left and right brain hemispheres. As illustrated in Figure 3, we first perform an affine registration of the input MRI scan to the MNI105 brain template [24]. This registration aims to unify the coordinate systems across different MR scans easing the learning and prediction of implicit surfaces.

Second, we construct a Cartesian grid of points at the desired resolution by dividing the template bounding box into evenly spaced points. Using the trained neural network $f_{\theta}$, we predict implicit surface representations at these points, represented by their continuous coordinates, for those four cortical surfaces according to the input MRI. Throughout this paper we reconstruct cortical surfaces using $512^{3}$ points, unless mentioned, providing sub-voxel precision. The model only needs to extract the feature maps once and can process the points in parallel, generating surfaces at high resolution efficiently.

Third, since most of the applications of cortical surface reconstruction require surfaces with spherical topology (i.e., genus zero) as discussed in Section 1, we employ a lightweight topology correction algorithm to fix the topological defects caused by wrong predictions of the implicit surface representation. Specifically, we use the method proposed by Bazin and Pham [3] which consists of a spherical level-set evolution that avoids evolving over critical points [30], outputting a 3D implicit surface volume with guaranteed spherical topology. Each target surfaces’ topology is corrected independently, and these tasks are performed in parallel for efficiency. Note that the proposed framework is agnostic to the topology correction method allowing the application of other techniques.

Finally, in order to obtain the mesh representation of the target surfaces, we apply a topology preserving marching cubes algorithm [20]. The predicted surfaces are in the template brain coordinate system, but these can be mapped back to the native space using the inverse of the transformation obtained during the registration of the input MR scan.

We now perform experiments with variations of the proposed DeepCSR to measure the importance of its main components. The quantitative results are presented in Table 1, while the qualitative results are shown in Figure 4. See below a brief summary of the dataset and evaluation metrics used, in addition to the detailed description of these experiments and the discussion of their results.

We now compare the precision, accuracy, and runtime of the proposed model to the FreeSurfer V6 (cross-sectional pipeline) and FastSurfer. More specifically, we use the DeepCSR (SDF) trained on the ADNI dataset and reconstruct surfaces at $512^{3}$ resolution as explained Section 3.3. Table 2 summarizes the results and Figure 5 shows examples of reconstructed surfaces.

Moreover, profiling DeepCSR, we observe that the initial registration takes $2.334~{}(\pm 0.022)$ minutes, the implicit surface prediction takes $8.672~{}(\pm 0.655)$ minutes, the topology correction takes $16.493~{}(\pm 1.064)$ minutes, and the marching cubes takes $0.325~{}(\pm 0.025)$ minutes. Therefore, further speed-up can be achieved by employing a faster topology correction method.