Introduction

Understanding the principles guiding neuronal organization has been a major goal in neuroscience. The ability to reconstruct individual neuronal arbors is necessary, but not sufficient to achieve this goal: understanding how neurons of the same and different types co-locate themselves requires the reconstruction of the arbors of multiple neurons sharing similar molecular and/or physiological features from the same brain. Such denser reconstructions may allow the field to answer some of the fundamental questions of neuroanatomy: do cells of the same type tile across the lateral dimensions by avoiding each other? To what extent do the organizational principles within a brain region extend across the whole brain? While dense reconstruction of electron microscopy images provides a solution (Denk and Horstmann 2004, Helmstaedter et al. 2013), its field-of-view has been limited for studying region-wide and brain-wide organization.

Recent advances in tissue clearing (Chung et al. 2013, Susaki et al. 2014) and light microscopy enable a fast, and versatile approach to this problem. In particular, oblique light-sheet microscopy can image thousands of individual neurons at once from the entire mouse brain at a 0.406 × 0.406 × 2.5 μm³ resolution (Narasimhan et al. 2017). Moreover, by registering reconstructed neurons from multiple brains of different neuronal gene expressions to a common coordinate framework such as the Allen Mouse Brain Atlas (Lein et al. 2007), it is possible to study neuronal structure and organization across many brain regions and neuronal cell classes. Therefore, this method may soon produce hundreds of full brain images, each containing hundreds of sparsely labeled neurons. However, scaling neuronal reconstructions to such large sets is not trivial. The gold standard of manual reconstruction is a tedious and labor-intensive process with a single neuronal reconstruction taking a few hours. This makes automated reconstruction the most viable alternative. Recently, many automated methods appeared for the reconstruction of neurons from light microscopy images. These include methods based on supervised learning with neuronal networks as well as other approaches (Peng et al. 2017, Türetken et al. 2011, Wang et al. 2011, Türetken, Benmansour, and Fua 2012, Sümbül et al. 2014, Gala et al. 2014). Some common problems include slow training and/or reconstruction speeds, tendency for topological mistakes despite high voxel-wise accuracy, and vulnerability to rare but important imaging artifacts such as stitching misalignments and microscope stage jumps. Here, we propose a supervised learning method based on a convolutional neural network architecture to address these shortcomings. In particular, we suggest (i) an objective function that penalizes topological errors more heavily, (ii) a data augmentation framework to increase robustness against multiple imaging artifacts, and (iii) a distributed scheme for scalability. Training data augmentation for addressing microscopy image defects was initially demonstrated for automated tracing of neurons in electron microscopy images (Lee et al. 2017). Here, we adapt this approach to sparse light microscopy images.

The U-Net architecture (Ronneberger, Fischer, and Brox 2015, Çiçek et al. 2016) has recently received significant interest, especially in the analysis of biomedical images. By segmenting all the voxels of an input patch rather than a central portion of it, the U-Net can learn robust segmentation rules faster, and decreases the memory and storage requirements. In this paper, we train a 3D U-Net convolutional network on a set of manually traced neuronal arbors. To overcome challenges caused by artifacts producing apparent discontinuities in the arbors, we propose a fast, connectivity-based regularization technique. While approaches that increase topological consistency exist (Briggman et al. 2009, Jain et al. 2010), they are either too slow for peta-scale images, or are not part of an online training procedure. Our approach is a simple, differentiable modification of the cost function, and the computational overhead scales linearly with the voxel count of the input patch. On the other hand, while these regularization techniques can enforce proper connectivity, there are relatively few examples of the various imaging artifacts in the training set. In order to increase the examples of such artifacts, we simulate them through various data augmentations and present these simulations under a unified framework. Taken together, our approach produces a significant increase in the topological accuracy of neuronal reconstructions on a test set.

In addition to accuracy, an efficient, scalable implementation is necessary for reconstructing petavoxel-sized image datasets. We maintain scalability and increase the throughput by using a distributed framework for reconstructing neurons from brain images, in which the computation can be distributed across multiple GPU instances. Finally, we augment data at run-time to avoid memory issues and computational bottlenecks. This significantly increases the throughput rate because data transfers are a substantial bottleneck. We report segmentation speeds exceeding 300 gigavoxels per hour and linear speedups in the presence of additional GPUs.

Methods

Convolutional neural network regularization through digital topology techniques

To create the training set, we obtain volumetric reconstructions of the manual arbor traces of neuronal images by a topology-preserving inflation of the traces (Sümbül et al. 2014). We use a 3D U-Net convolutional neural network architecture (Ronneberger, Fischer, and Brox 2015, Çiçek et al. 2016, Lee et al. 2017) to learn to segment the neurons from this volumetric training set. Since neuronal morphology is ultimately represented and analyzed as a tree structure, we consider the branching pattern of the segmented neuron more important than its voxelwise accuracy. Hence, to penalize topological changes between the ground-truth and the prediction at the time of training, we binarize the network output by thresholding and identify all non-simple points in this binarized patch based on 26-connectivity (Bertrand and Malandain 1994) — points when added or removed change an object’s topology (e.g., splits and mergers) — and assign larger weights to them in the binary cross-entropy cost function
$$\label{equ:loss} % J(\hat{y}, y) = -\frac{1}{N}\sum^{N}_{i=1} w_{i} \big[ y_{i} \log % (\sigma(\hat{y}_{i})) + (1 - y_{i}) \log (1 - {\sigma(\hat{y}_{i})}) % \big] J(\hat{y}, y) = -\frac{1}{N}\sum^{N}_{i=1} w_{i} \big[ y_{i} \log (\hat{y}_{i}) + (1 - y_{i}) \log (1 - {\hat{y}_{i}}) \big]$$
where w_i = w > 1 if voxel i is non-simple while w_i = 1 otherwise, N is the number of voxels, and y_i and $\hat{y_i}$ are the label image and predicted segmentation, respectively. Note that the simple-ness of a voxel depends only on its 26-neighborhood, and therefore this operation scales linearly with the patch size.

Simulation of image artifacts through data augmentations

Data augmentation is a technique that augments the base training data with pre-defined transformations of it. By creating statistical invariances (e.g. against rotation) within the dataset or over-representing rarely occurring artifacts, augmentation can increase the robustness of the learned algorithm. Motivated by the fact that 3D microscopy is prone to several image artifacts, we followed a unified framework for data augmentation. In particular, our formalism requires explicit models of the underlying artifacts and the desired reconstruction in their presence to augment the original training set with simulations of these artifacts.

We define the class of “artifact-generating” transformations as S such that if 𝒯 ∈ S, then 𝒯 = 𝒯_R ⊗ 𝒯_L for 𝒯_R : ℝ^{n₁ × n₂ × n₃} → ℝ^{n₁ × n₂ × n₃} and 𝒯_L : {0, 1}^{n₁ × n₂ × n₃} → {0, 1}^{n₁ × n₂ × n₃}, where 𝒯_R acts on an n₁ × n₂ × n₃ raw image and 𝒯_L acts on its corresponding label image. For example, the common augmentation step of rotation by 90^∘ can be realized by 𝒯_R and 𝒯_L both rotating their arguments by 90^∘. Data augmentation adds these rotated raw/label image pairs to the original training set (Fig. 1).

Occluded branches: Branch occlusions can be caused by photobleaching or an absence of a fluorophore. We model the artifact-generating transformation for an absence of a fluorophore as 𝒯 = 𝒯_R ⊗ ℐ, where
𝒯_ℛ(R; x, y, z) = R − PSF(x, y, z)
such that ℐ denotes the identity transformation, x denotes the position of the absent fluorophore and PSF is its corresponding point-spread function. Here, we approximated the PSF of a fluorophore with a multivariate Gaussian.

Duplicate sections: The stage of a scanning 3D microscope can intermittently stall, which can duplicate the imaging of a tissue section. The artifact-generating transformation for stage stalling is given by 𝒯 = 𝒯_R ⊗ ℐ, where
$$\mathcal{T_R}(R; \mathbf{r}, \mathbf{r_0}, N) = \begin{cases} R(\mathbf{r}), \mathbf{r} \not\in N \\ R(\mathbf{r_{0}}), \mathbf{r} \in N \end{cases}$$
for the region r = (x, y, z) and the plane r₀ = (x₀, y, z) such that 𝒯_R duplicates the slice r₀ in a rectangular neighborhood N.

Dropped sections: Similar to the stalling of the stage, jumps that result in missed sections can occur intermittently. The corresponding artifact-generating transformation is given by 𝒯 = 𝒯_R ⊗ 𝒯_L, where
$$\mathcal{T}_{R}(R; \mathbf{r}, x_0, \Delta) = \begin{cases} R(x, y, z), & x \leq x_0 \\ R(x + \Delta, y, z), & x > x_0 \end{cases}$$
and
$$\mathcal{T}_{L}(L; \mathbf{r}, x_0, \Delta) = \begin{cases} L(x, y, z), & x \leq x_0 - \Delta \\ L(D(x), y, z), & |x-x_0-\frac{\Delta}{2}| > \frac{3\Delta}{2} \\ L(x + 2\Delta, y, z), & x \geq x_0 + 2\Delta \end{cases}$$
such that r = (x, y, z), for D(x, x₀, Δ) = x₀ − Δ + ³/₂⌈x − x₀ + Δ⌉, which downsamples the region to maintain partial connectivity in the label. Hence, 𝒯_R skips a small region given by Δ at x₀, and 𝒯_L is the corresponding desired transformation on the label image.

Stitching misalignment: Misalignments can occur between 3D image stacks, potentially causing topological breaks and mergers between neuronal branches. The corresponding artifact-generating transformation is given by 𝒯 = 𝒯_R ⊗ 𝒯_L, where
$$\mathcal{T}_{R}(R;x,y,z,\Delta) = \begin{cases} R(x, y, z), & x \leq x_{0} \\ R(x, y + \Delta, z), & x > x_{0} \end{cases}$$
and
$$\mathcal{T}_{L}(L;x,y,z,\Delta) = \begin{cases} L, & x \leq x_{0} - \frac{1}{2} \Delta \\ \Sigma_{zy}(\Delta) L, & |x-x_{0}|<\frac{1}{2} \Delta \\ L(x, y + \Delta, z), & x > x_{0} + \frac{1}{2} \Delta \end{cases}$$
such that Σ_zy(Δ) is a shear transform on L. Hence, 𝒯_R translates a region of R to simulate a stitching misalignment, and 𝒯_L shears a region around the discontinuity to maintain 18-connectivity in the label.

Light scattering: Light scattering by the cleared tissue can create an inhomogeneous intensity profile and blur the image. To simulate this image artifact, we assumed the scatter has a homogeneous profile and is anisotropic due to the oblique light-sheet. We approximate these characteristics with a Gaussian kernel: G(x, y, z) = G(r) = 𝒩(r; μ, Σ). In addition, the global inhomogeneous intensity profile was simulated with an additive constant. Thus, the corresponding artifact-generating transformation is given by 𝒯 = 𝒯_R ⊗ ℐ, where
𝒯_R(R) = R(x, y, z) * G(x, y, z) + λ

Fully automated, scalable tracing

To optimize the pipeline for scalability, we store images as parcellated HDF5 datasets. For training, a file server software streams these images to the GPU server, which performs data augmentations on-the-fly, to minimize storage space requirements. For deploying the trained neural network, the file server similarly streams the datasets to a GPU server for segmentation. Once the segmentation is completed, the neuronal morphology is reconstructed automatically from the segmented image using the UltraTracer neuron tracing tool within the Vaa3D software package (Peng et al. 2017).

Experimental Procedure

In our experiments, we used a dataset of 54 manually traced neurons imaged using oblique light-sheet microscopy. These morphological annotations were dilated while preserving topology for training the neural network for segmentation. We partitioned the dataset into training, validation, and test sets by randomly choosing 25, 8, and 21 neurons, respectively. The software package PyTorch was used to implement the neural network (Paszke et al. 2017). The network was trained using an Adam optimizer for gradient descent (Kingma and Ba 2014). Training and reconstruction were conducted on two Intel Xeon Silver 4116 CPU, 256 GB RAM, and 2 NVIDIA GeForce GTX 1080 Ti GPUs.

Results

Topologically accurate reconstruction

To quantify the topological accuracy of the network on light-sheet microscopy data, we define the topological error as the number of non-simple points that must be added or removed from a prediction to obtain its corresponding label. Specifically, for binary images L̂ and L, let 𝒲(L̂, L) denote a topology-preserving warping of L̂ that minimizes the voxelwise disagreements between the warped image and L (Jain et al. 2010, Sümbül et al. 2014), L̂ ∩ L denote the binary image whose foreground is common to both L̂ and L, and c(L) denote the number of foreground voxels of L. We quantify the agreement between a reconstruction L̂ and label L using the Jaccard index as
$$J(\hat{L}, L) = \frac{c(\mathcal{W}(\hat{L}, L) \cap L)}{c(\mathcal{W}(\hat{L}, L) \cup L)}.$$
We compared this score across different U-Net results: without any augmentations or regularization, with the augmentations, with the topological regularization, and with both the topological regularization and the augmentations. The U-Net results with augmentations and topological regularization performed significantly better compared to the results without augmentations or regularization (Figs 2, 3).

Neuron reconstruction is efficient and scalable

To quantify the efficiency of the distributed framework, we measured the framework’s throughput for augmenting data, training on the data, and segmenting the data. Augmentations performed at 35.2 ± 9.2 gigavoxels per hour while training performed at 16.8 ± 0.2 megavoxels per hour. Segmentation performed at 348.8 ± 1.9 gigavoxels per hour. Both segmentation and training showed a linear speedup with an additional GPU. For an entire mouse brain, neuronal reconstruction would take about 23 hours on a single GPU.

Discussion

In this paper, we proposed an efficient, scalable, and accurate algorithm capable of reconstructing neuronal anatomy from light microscopy images of the whole brain. Our method employs topological regularization as well as simulates discontinuous image artifacts inherent to the imaging systems. These techniques help maintain topological correctness of the trace (skeleton) representations of neuronal arbors.

While we demonstrated the merit of our approach on neuronal images obtained by oblique light-sheet microscopy, our methods address some of the problems common to most 3D fluorescence microscopy techniques. Therefore, we hope that some of our methods will be useful for multiple applications. Combined with the speed and precision of oblique light-sheet microscopy, the distributed and fast nature of our approach enables the production of a comprehensive database of neuronal anatomy across many brain regions and cell classes. We believe that these aspects will be useful in discovering different cortical cell types as well as understanding the anatomical organization of the brain.