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
π Todo List
- Multi-channel segmentation.
- Generalization to other microscopes. E.g., (ExASPIM, TissueCyte)
Inference
It supports two processing modes:
- Lazy segmentation: For large datasets that cannot fit into memory, images are resized to 1024x1024, normalized, segmented, upsampled, and stored in a Zarr format.
- 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 .