Reconstructing neuronal anatomy from whole-brain images

TL;DR

This paper introduces a scalable, automated method for reconstructing neuronal anatomy from whole-brain light microscopy images, addressing artifacts and discontinuities with neural network-based techniques.

Contribution

It presents connectivity-preserving neural network methods and a scalable pipeline for automated, high-resolution neuronal reconstruction from whole-brain images.

Findings

Effective neural network-based reconstruction pipeline

Handling of image artifacts and discontinuities

Scalable implementation for large datasets

Abstract

Reconstructing multiple molecularly defined neurons from individual brains and across multiple brain regions can reveal organizational principles of the nervous system. However, high resolution imaging of the whole brain is a technically challenging and slow process. Recently, oblique light sheet microscopy has emerged as a rapid imaging method that can provide whole brain fluorescence microscopy at a voxel size of 0.4 by 0.4 by 2.5 cubic microns. On the other hand, complex image artifacts due to whole-brain coverage produce apparent discontinuities in neuronal arbors. Here, we present connectivity-preserving methods and data augmentation strategies for supervised learning of neuroanatomy from light microscopy using neural networks. We quantify the merit of our approach by implementing an end-to-end automated tracing pipeline. Lastly, we demonstrate a scalable, distributed implementation that can reconstruct the large datasets that sub-micron whole-brain images produce.

Figures

Equations

$$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{]}$$

$${\mathcal{T}}_{\mathcal{R}}(R;x,y,z)=R-\mathrm{PSF}(x,y,z)$$

$${\mathcal{T}}_{\mathcal{R}}(R;\mathbf{r},{\mathbf{r}}_0,N)=\{\begin{array}{ll}R(\mathbf{r}),& \mathbf{r}\notin N\\ R({\mathbf{r}}_0),& \mathbf{r}\in N\end{array}$$

$${\mathcal{T}}_R(R;\mathbf{r},x_0,\Delta )=\{\begin{array}{ll}R(x,y,z),& x\le x_0\\ R(x+\Delta ,y,z),& x>x_0\end{array}$$

$${\mathcal{T}}_L(L;\mathbf{r},x_0,\Delta )=\{\begin{array}{ll}L(x,y,z),& x\le 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\ge x_0+2\Delta \end{array}$$

$${\mathcal{T}}_R(R;x,y,z,\Delta )=\{\begin{array}{ll}R(x,y,z),& x\le x_0\\ R(x,y+\Delta ,z),& x>x_0\end{array}$$

$${\mathcal{T}}_L(L;x,y,z,\Delta )=\{\begin{array}{ll}L,& x\le 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{array}$$

$${\mathcal{T}}_R(R)=R(x,y,z)\ast G(x,y,z)+\lambda$$

$$J(\hat{L},L)=\frac{c(\mathcal{W}(\hat{L},L)\cap L)}{c(\mathcal{W}(\hat{L},L)\cup L)}.$$

Full text

00footnotetext: ⋆ Corresponding author: James Gornet, james.gornet@columbia.edu. This research is supported through a grant from the National Institutes of Health (NIMH U01MH114824). Work performed at the Allen Institute.

**Index Terms— ** image segmentation, light microscopy, machine learning

1 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 [1, 2], its field-of-view has been limited for studying region-wide and brain-wide organization.

Recent advances in tissue clearing [3, 4] 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 $\times$ 0.406 $\times$ $2.5\text{\,}\mathrm{\SIUnitSymbolMicro}\mathrm{m}^{3}$ resolution [5]. 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 [6], 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 [7, 8, 9, 10, 11, 12]. 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 [13]. Here, we adapt this approach to sparse light microscopy images.

The U-Net architecture [14, 15] 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 [16, 17], 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.

2 Methods

2.1 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 [18]. We use a 3D U-Net convolutional neural network architecture [14, 15, 13] 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 [19] — 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

[TABLE]

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.

2.2 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 $\mathcal{T}\in S$, then $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{T}_{L}$ for $\mathcal{T}_{R}:\mathbb{R}^{n_{1}\times n_{2}\times n_{3}}\rightarrow\mathbb{R}^{n_{1}\times n_{2}\times n_{3}}$ and $\mathcal{T}_{L}:{\{0,1\}}^{n_{1}\times n_{2}\times n_{3}}\rightarrow{\{0,1\}}^{n_{1}\times n_{2}\times n_{3}}$, where $\mathcal{T}_{R}$ acts on an $n_{1}\times n_{2}\times n_{3}$ raw image and $\mathcal{T}_{L}$ acts on its corresponding label image. For example, the common augmentation step of rotation by $90^{\circ}$ can be realized by $\mathcal{T}_{R}$ and $\mathcal{T}_{L}$ both rotating their arguments by $90^{\circ}$. 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 $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{I}$, where

[TABLE]

such that $\mathcal{I}$ denotes the identity transformation, $x$ denotes the position of the absent fluorophore and $\mathrm{PSF}$ is its corresponding point-spread function. Here, we approximated the $\mathrm{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 $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{I}$, where

[TABLE]

for the region $\mathbf{r}=(x,y,z)$ and the plane $\mathbf{r_{0}}=(x_{0},y,z)$ such that $\mathcal{T}_{R}$ duplicates the slice $\mathbf{r_{0}}$ 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 $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{T}_{L}$, where

[TABLE]

and

[TABLE]

such that $\mathbf{r}=(x,y,z)$, for $D(x,x_{0},\Delta)=x_{0}-\Delta+\frac{3}{2}\lceil x-x_{0}+\Delta\rceil$, which downsamples the region to maintain partial connectivity in the label. Hence, $\mathcal{T}_{R}$ skips a small region given by $\Delta$ at $x_{0}$, and $\mathcal{T}_{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 $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{T}_{L}$, where

[TABLE]

and

[TABLE]

such that $\Sigma_{zy}(\Delta)$ is a shear transform on $L$. Hence, $\mathcal{T}_{R}$ translates a region of $R$ to simulate a stitching misalignment, and $\mathcal{T}_{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(\mathbf{r})=\mathcal{N}(\mathbf{r};\mu,\Sigma)$. In addition, the global inhomogeneous intensity profile was simulated with an additive constant. Thus, the corresponding artifact-generating transformation is given by $\mathcal{T}=\mathcal{T}_{R}\otimes\mathcal{I}$, where

[TABLE]

2.3 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 [7].

3 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 [20]. The network was trained using an Adam optimizer for gradient descent [21]. Training and reconstruction were conducted on two Intel Xeon Silver 4116 CPU, 256 GB RAM, and 2 NVIDIA GeForce GTX 1080 Ti GPUs.

4 Results

4.1 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 $\hat{L}$ and $L$, let $\mathcal{W}(\hat{L},L)$ denote a topology-preserving warping of $\hat{L}$ that minimizes the voxelwise disagreements between the warped image and $L$ [17, 11], $\hat{L}\cap L$ denote the binary image whose foreground is common to both $\hat{L}$ and $L$, and $c(L)$ denote the number of foreground voxels of $L$. We quantify the agreement between a reconstruction $\hat{L}$ and label $L$ using the Jaccard index as

[TABLE]

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).

4.2 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 $\pm$ 9.2 gigavoxels per hour while training performed at 16.8 $\pm$ 0.2 megavoxels per hour. Segmentation performed at 348.8 $\pm$ 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.

5 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.