README.md

---
license: mit
language:
- en
tags:
- segmentation
- microscopy
- tissue
---
# aind-brain-tissue-segmentation

<p align="center">
    <img src="https://allenneuraldynamics.github.io/assets/img/AIND_logo.png" width="400"/>
<p>

<p align="center">
    <a href=""><b>AIND</b></a> &nbsp&nbsp ｜ &nbsp&nbsp 🖥️ <a href="https://github.com/AllenNeuralDynamics/aind-brain-segmentation">GitHub</a> &nbsp&nbsp  | &nbsp&nbsp🤗 <a href="https://huggingface.co/AllenInstitute">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤖 <a href="https://www.allenneuraldynamics.org/">Allen Institute for Neural Dynamics</a>&nbsp&nbsp
<br>

-----

[**Brain Tissue Segmentation - Lightsheet Microscopy**]() <be>

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

<a href="https://github.com/user-attachments/assets/1021a3a3-cfa1-460c-b1d7-29746ecf764c">Video Sample 771602</a>

<a href="https://github.com/user-attachments/assets/66ec0a36-6798-4081-98e4-b6e744355e99">Video Sample 782499</a>

## 📑 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**.

<a href="https://github.com/AllenNeuralDynamics/aind-brain-segmentation/blob/main/code/large_scale_mask_prediction.py">Whole-brain segmentation script</a>


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