User:Ahu Qulab/sandbox

mBrianAligner is a cross-modality image registration pipeline to support whole brain mapping projects. In addition to aligning 3D mouse brain images of different modalities, mBrainAligner also enables mapping digitally reconstructed compartments (e.g., dendritic, axonal, and soma distributions) to a common target space to faciliate the visualization, comparison and analysis. mBrainAligner is not limited to the use of intra- and cross-modality mouse brain registration, it may also be extended to partially imaged data and multi-timepoint registration or other species.

Recent whole brain mapping projects are collecting large-scale 3D images using powerful and informative modalities, such as STPT, fMOST, VISoR, LSFM or MRI. Registration of these multi-dimensional whole-brain images onto a standard atlas is essential for characterizing neuron types and constructing brain wiring diagrams. However, cross-modality image registration is challenging due to intrinsic variations of brain anatomy and artifacts resulted from different sample preparation methods and imaging modalities. mBrainAligner is precisely to break down the barriers between these modal data so that they can be comprehensively analyzed in a unified common space.mBrainAligner can not only complete tasks such as brain region division, soma cell localization and axon projection prediction, but also has good adaptability to cross-modal scenes and has excellent cross-registration performance between different modal brain images.

mBrainAligner contains three modules: (1) image preprocessing and global registration, (2) Coherent Landmark Mapping (CLM) based automatic registration, and (3) optional semi-automatic refinement. To accommodate different registration accuracy or throughput requirements, the above modules can be concatenated or executed separately. In addition, some useful tools including 2.5D corner detector, stripe artifacts removal, and image or metadata warping tools are also provided in this package for the user's convenience.

Application background
Recent advances in high-resolution light microscopy, tissue clearing, sparse labeling techniques, and neuron tracing methods now make it feasible to map the whole mammalian brain at single-cell resolution. Large international efforts such as the BRAIN Initiative Cell Census Network (BICCN), MouseLight project, Allen Mouse Brain Connectivity Atlas 25 and Mouse Connectome project, are engaged in cell typing, mapping long-range axonal projections and microcircuit connection analyses. Massive image datasets are acquired through a variety of high-resolution, high-throughput imaging techniques, such as serial two-photon tomography (STPT), fluorescence micro-optical sectioning tomography (fMOST), light-sheet fluorescence microscopy (LSFM), or volumetric imaging with synchronous on-the-fly-scan and readout (VISoR).

A critical technique required to use these data is registration (alignment) of brains, their compartments, and various neuronal structures including somas, dendrites, and axon arbors. Data that are acquired across subjects, modalities, time or even species must be mapped onto a canonical coordinate space. Its importance is obvious in multiple ways. First, registration of datasets of different modalities generated in different projects provides a valuable resource for ever-expanding studies of neuron wiring and functions. Second, the registration to a common standard atlas enables digital delineation and quantification of brain anatomy and neuron morphology and other attributes such as transcriptomics. Third, registration allows quantitative comparison of features of neurons across samples under different conditions, thus enabling analysis of distinct aspects of neurobiology in a coordinated manner.

Established brain atlases such as the Allen Common Coordinate Framework atlas (CCFv3) (using STPT imaging) 29,30 and its derived LSFM domain atlas 31 have helped neurobiologists to study gene expression patterns, brain anatomy, and neural circuits. Currently, fMOST imaging is increasingly used for whole-brain single cell analysis. High-resolution images in this modality are being generated (e.g. in the BICCN program) to produce large single-neuron morphology databases 32. It is crucial to register such data to the Allen brain atlas to compare mapped neurons and populations of neurons generated using different experimental conditions and modalities.

Removal of stripe artifacts
The stripe noise present in the raw fMOST images is mainly caused by fluorescent bleaching during the knife cutting and imaging process. We modeled this kind of stripe artifacts as multiplicative periodic noise and designed a log-space frequency notch filter to realize high-quality stripe artifacts removal. For each fMOST image, (1) one 2D coronal slice with the largest foreground brain area was extracted (Extended Data Fig.a) and converted to the frequency domain using fast Fourier transform (Extended Data Fig.b). (2) By examining the frequency spectrum, the bandwidth, orientation and cutoff frequency of the notch filter were identified and then used to construct a Gaussian notch filter (Extended Data Fig.c). (3) We imposed the log transform on the intensity of the fMOST image in the spatial domain to convert the multiplicative noise to additive one (Extended Data Fig.d). (4) The log-transformed fMOST image was filtered in the frequency domain in a slice-by-slice manner using the Gaussian notch filter generated in step 2 (Extended Data Fig.e). (5) Finally, we computed the inverse fast Fourier transform on the filtered spectrum and used inverse-log transform to bring the filtered image back from the log space (Extended Data Fig.f).

Robust point matching-based automatic global registration
The function of global registration is to make the image to be registered have the same size and direction as the target image, which will help reduce the difficulty of local registration in the later stage and improve the registration accuracy.

To minimize the impact of intensity variation between different modalities, we implement the automatic global alignment by affine aligning brains’ outer-contour points through point-set matching. First, the foreground brain region was extracted via Otsu thresholding53, followed by morphology filtering to remove the small and thin artifacts. Then we uniformly sampled ( grid) the contour of the foreground region to obtain a point-set (~1500 points per brain) that represents the shape of brain’s outer-contour. To tackle the large direction variation between brains, the principal axes of target and subject point-sets were extracted using principal component analysis, and then rigidly aligned. Finally, following the deterministic annealing framework of robust point matching54, we iteratively refine an affine transform to maximize the overall shape of two point-sets. We set the initial annealing temperature to 0.25, and gradually decrease the temperature with an annealing rate of 0.96 during iteration.



CLM-based automatic local registration
A key step of mBrainAligner is stable and accurate nonlinear local registration, which is achieved through two graphs The correspondence between dense points on the image maps the image to be registered to the target image. The CLM algorithm performs correlation deformation on the points on the target brain to match the brain to be registered, and makes the points on the target imageand adjacent points can find the best matching position on the image to be registered in a coordinated manner. Since the points before and after deformation naturally form matching point pairs, CLM essentially avoids the unreliable detection and identification of error-prone points. Cross-modal point matching. To ensure the robustness and coherence of point deformations, the deformation process is driven by three complementary features (gradients, oriented gradient histograms, and segmentation probabilities), which encode contours, shape context features, and semantic information respectively, and are embedded by CCFv3 (or fMOST atlas) prior features of the brain anatomy shape for normalization constraints. In addition, a DNN model is used to estimate the brain region segmentation probability map. Since brain region segmentation is an important application based on brain atlas registration, the registration task can also benefit from the segmentation results; the semantic information in the segmentation can be incorporated as discriminative features for registration use. Using the CLM algorithm we work together to ensure the robustness and accuracy of mBrainAligner in challenging cross-modal scenarios by integrating coherence mapping strategies, discriminative feature integration, and the utilization of effective shape priors.

Coherent landmark mapping
CLM establishes a dense correspondence between two images by coherently deforming target landmarks to fit the subject brain. The landmarks before and after deforming naturally form matching pairs, and were used to calculate the displacement feld for local image warping.

Suppose we have a subject image that consists of $$ N $$ voxels $$ V = \{v_i, i = 1, 2, \ldots, N\} $$, and a target landmark-set of $$ M $$ points $$ L = \{l_j, j = 1, 2, \ldots, M\} $$, where $$ v_i $$ is the $$ i_{th} $$ voxel and $$ l_j $$ is the $$ j_{th} $$  landmark in subject and target image, respectively. Let $$ P = \{p_{ij}, p_{ij} \in [0, 1]\} $$ be the classification probability of $$ i_{th} $$ voxel to the $$ j_{th} $$ landmark, we form the CLM as the following energy minimization problem: $$ 	E(P,\theta) = \sum_{j=1}^{M} \sum_{i=1}^{N} p_{ij}(1-s_{ij}) \frac{||v_i - f(l_j,\theta)|| ^2}{r^2}+ T\sum_{j=1}^{M} \sum_{i=1}^{N} p_{ij} log p_{ij}, s.t.\sum_{i=1}^{N}p_{ij} =1 $$ where $$ f(l_j,\theta) $$ is a non-linear spatial transform which maps the landmark $$ l_j $$ to a new location with parameters $$ \theta $$, $$ s_{ij} \in [0, 1] $$ denotes the matching score between voxel and currently deformed landmark $$ f(l_j,\theta) $$, and $$ r $$ is the distance influence factor which controls the impact of voxel-to-landmark distance to their classification probability. Minimizing the first term encourages assigning larger classification probability to the voxel-to-landmark pair with higher matching score and smaller Euclidian distance. While minimizing the second term is equivalent to maximizing the entropy of $$ P $$, which aims to enforce the convexity of the energy function and alleviate the risk of algorithm getting trapped in the local minimum during optimization. $$ T $$ acts as a balance factor between two terms.

We adopted a dual update strategy to solve the joint classification and spatial transform optimization problem posed in energy function. In each iteration, we first fix the transformation $$ \theta $$ and calculate the voxel-to-landmark classification with currently deformed landmarks, then update the transform and perform landmark deformation with classification probability fixed. We denote the above two steps as landmark-guided voxel classification and voxel classification-guided landmark deformation respectively.

Landmark-guided voxel classification
Given spatial transform $$ \theta $$, the deformed target landmarks can be obtained as $$ f(L,\theta) $$. For any $$ i $$, $$ j $$ value, minimizing $$ E(P,\theta) $$ with respect to $$ P $$  will lead to: $$ 	\frac{\partial E}{\partial p_{ij}} = 0 \Rightarrow p_{ij} = exp \left (-\frac { (1-s_{ij}) ||v_i - f(l_j,\theta)||^2}{r^2 T}-1 \right) $$ we choose $$ T=1 $$, and normalize classification probability so that $$ \sum_{i}p_{ij} =1 $$ for any $$ i \in [0, N] $$ once they are updated using equation. To speed up the iterative voxel classification, we only concern a subdomain(21×21×21) around each landmark and choose $$ r=10 $$.

Voxel classification-guided landmark deformation
With fixed voxel classification probability $$ P $$, we realized the coherent landmark deforming in three steps. First, we preliminarily updated each landmark to the weighted center of mass of all voxels $$ l_{j}^{pre}= \sum_{i=1}^{N}p_{ij}v_{i} $$. Then we updated transform parameters $$ \theta $$ so that deformed landmarks $$ f(L,\theta) $$ can best approximate their preliminary updated ones $$ \{l_j^{pre}, j = 1, 2, \ldots, M\} $$ under smoothness and brain shape prior constraints. We choose smooth-thin-plate-spline (STPS) to model the spatial transform with parameters $$ \theta $$, which can be decomposed into an affine transform and a smoothness parameter $$ \lambda $$ regularized non-linear warping component. If $$ \lambda = 0 $$, an exact interpolation was performed. We choose $$ \lambda = 0.001 $$ to balance the goodness of fit between two point sets and the smoothness of deformation. The spatial transform parameters $$ \theta $$ were solved as a least-squares solution of STPS. Finally, we used LittleQuickWarp (a fast and memory efficient implementation of STPS) with updated $$ \theta $$ to transform the initial target landmarks and realize the coherent landmark deforming. Since target landmarks are obtained on CCFv3 annotation template, we ensure that all landmarks are on the boundaries of brain region. They naturally capture the shape prior of these regions and form a simplified point-set representation of the CCFv3 atlas. The smooth (regularized by STPS) and coherent deformation of such a simplified atlas onto the subject brain implicitly utilizes the shape priors and spatial relationship of brain regions, while these priors are hard to be incorporated in the commonly used B-spline based methods.

If provided, the modality-specific shape priors can also be accommodated under the CLM framework to further enhance the robustness and accuracy of registration. Given the fMOST atlas, since the transformation between CCFv3 and fMOST atlas was known and fixed, we utilized these modality-specific shape priors by first mapping the initial landmarks of different brain regions to fMOST space, and then deforming them in a region-wise manner following the above mentioned STPS deforming strategy. In our implementation, this region-wise constraint was applied before the brain-wise constraint in each iteration.

Matching score computation
We based the calculation of matching score between a voxel-to-landmark pair upon three complementary features: (i) gradient, (ii) histogram of oriented gradients (HOG), and (iii) segmentation probability. Since these features encode the edge, shape context and semantic information respectively, their integration is expected to provide CLM with more reliable matching metric to better cope with the challenges of cross-modal registration.

gradient feature
We generated the gradient feature by applying the Sobel filter (3×3×3) on Gaussian smoothed (3×3×3) volume.

Hog feature
At a given location, its HOG feature was calculated by concatenating the normalized gradient histograms of eight blocks within the patch (21×21×21) that centered on that location. We divided each patch into 27 cells (7×7×7) and combined its neighboring (2×2×2) cells to create blocks. For each block, the gradient histogram (nine orientations) of each cell within it was computed, concatenated and normalized to form the block feature. This resulted in a 576-dimensional HOG feature for each location.

segmentation probability
We choose to generate the segmentation probability feature using a semantic segmentation network. As a proof of principle, we adopted a slightly modified 3D U-Net59 to generate the segmentation probability of each voxel to six main brain regions (HY, HPF, CTX, CBX, BS, CP) and background. Indeed, under the CLM framework, 3D U-Net can be readily substituted with other more sophisticated semantic segmentation networks to further improve the registration performance. The network architecture used in this study is shown in Extended Data Fig. a, which contains a classical encoder-decoder structure. We use four layers at the encoder stage and four layers at the decoder stage, and they have 32, 64, 128 and 256 channels respectively. The 7-dimensional score maps generated at the final softmax layer were used as the segmentation probability feature. The network was trained from scratch on 57 brains with ground truth segmentation labels obtained via mBrainAligner’s semi-automatic registration module. We used 23 brains as the training set, 3 brains as the validation set and the remaining for testing. The segmentation results of six brain regions are visualized in Extended Data Fig. b. Potential segmentation errors can be mitigated as the segmentation information is only one of the three discriminative features that guide the deformation of landmarks, for which the deformation is further constrained by the overall shape priors.

similarity metrics
Finally, we computed the matching score of each landmark-to-voxel pair by accumulating the similarity score of their centered patches upon three feature maps. In our implementation, (i) inverse mean squared error, (ii) mutual information, and (ii) normalized cross-correlation were used as similarity metrics.

Semiautomatic refinement
The semi-automatic registration module provides an iteractive interface for investigators to visualize, examine and fine-tune the registration results (Extended Data Fig. a). In this interface, the boundary of brain regions defined in CCFv3 was vectorized as Bezier curves and visualized as a layer overlaid on the registered brain. The control nodes of Bezier curves were evenly sampled from brain region boundaries defined in CCFv3 annotation template. The registration quality could be examined by checking how well the registered brain would fit CCFv3 boundaries in coronal, sagittal and horizontal views at different scales. Different to reported method that performed slice-by-slice adjustment in coronal view60, our method could fine-tune three-view panels synchronously by simply dragging the curves, thus distortions in Z direction could also be rectified. In addition, we also provided a magnetic lasso61 mode to assist in the alignment of curves in a semi-automatic manner. Three levels of sub-region granularity (coarse, medium, fine) were provided for each main brain region to accommodate different local granularity requirements (Extended Data Fig. b). We realize the final 3D image deforming by using STPS transformation that computed upon the offsets of fine-tuned control nodes. We demonstrated that the semi-automatic registration module could register partially imaged brain, a known challenging registration problem for fully automatic methods.

How to use mBrainAligner
Detailed tutorial on using mBrainAligner[https://github.com/Vaa3D/vaa3d_tools/tree/master/hackathon/mBrainAligner)