aind-brain-tissue-segmentation

AIND    ｜    🖥️ GitHub    |   🤗 Hugging Face   |   🤖 Allen Institute for Neural Dynamics  

---

Brain Tissue Segmentation - Lightsheet Microscopy

In this repository, we present a tool for fast brain tissue segmentation in lightsheet microscopy. This model was trained in entire 2D slices of data coming from our SmartSPIM Lightsheet Microscopes. We used a downsampled resolution corresponding to 14.4 microns in XY and 16.0 microns in Z that then were resized to 1024x1024.

We also provide a script that perform fast inference in entire brains stored in OMEZarr format.

Brain Segmentation Examples

Video Sample 771602

Video Sample 782499

📑 Todo List

• Multi-channel segmentation.
• Generalization to other microscopes. E.g., (ExASPIM, TissueCyte)

Inference

It supports two processing modes:

1. Lazy segmentation: For large datasets that cannot fit into memory, images are resized to 1024x1024, normalized, segmented, upsampled, and stored in a Zarr format.
2. In-memory segmentation: For smaller datasets, a multiscale version is loaded into memory for processing.

The model was trained with a resolution of 14.4 microns in XY and 16 microns in Z.

Whole-brain segmentation script

python large_scale_mask_prediction.py --image_path path/to/your/image.zarr --model_path path/to/your/model.safetensors

Input Parameters

Parameter	Description
image_path	Path to the OME-Zarr dataset (local or S3). Example: s3://bucket/SmartSPIM_Dataset/Ex_639_Em_667.zarr
model_path	Path to the trained segmentation model. It can point to a .cpkt or .safetensors file.
output_folder	Directory where results will be saved.
target_size_mb	Memory allocation target for large-scale predictions.
n_workers	Number of workers for PyTorch DataLoader.
super_chunksize	Data shard size for optimized cloud communication.
scale	Name of the multiscale dataset to use for segmentation.
scratch_folder	Path for temporary files (Raw data).
image_height	Resize height for segmentation (Default: 1024).
image_width	Resize width for segmentation (Default: 1024).
prob_threshold	Probability threshold for segmentation mask generation.

Outputs

Output File	Description
probabilities.zarr	Zarr dataset containing segmentation probabilities (useful for post-processing).
segmentation_mask.zarr	Segmentation mask in the original image space (before resizing).
data.zarr	Raw input data used for segmentation (optional).

Installation

Clone the repo:

git clone https://github.com/AllenNeuralDynamics/aind-brain-segmentation
cd aind-brain-segmentation

Install dependencies:

docker build -t brain_seg -f Dockerfile .