Neuron Segmentation

A Python package for performing neuron segmentation in ExaSPIM image datasets, designed for large-scale, high-resolution volumetric data. The pipeline combines deep learning–based affinity prediction with graph-based algorithms to produce accurate neuron reconstructions.

Overview

The segmentation pipeline consists of three main steps:

1. Affinity Prediction: 3D UNet predicts affinities between neighboring voxels.

2. Generate Supervoxels: Seeded watershed transforms affinities into an initial supervoxel oversegmentation.

3. Agglomerate Supervoxels: Supervoxels are iteratively merged using a graph-based algorithm.

In addition, this repository provides tools for skeletonization and saving the result as a ZIP archive of SWC files.

Figure: Visualization of segmentation pipeline.

Installation

To use the software, install c++ dependencies:

sudo apt install libboost-dev

then run the following in the root directory:

pip install -e .

Predict

Here’s an example of running the segmentation pipeline.

from aind_exaspim_neuron_segmentation import inference
from aind_exaspim_neuron_segmentation.utils import img_util


# Parameters
affinity_mode = True
patch_shape = (96, 96, 96)

# Load model
model_path = "path-to-model-weights"
model = inference.load_model(model_path, affinity_mode=affinity_mode)

# Read image
img_path = "path-to-image"
img = img_util.read(img_path)

# Predict affinites
affinites = inference.predict(
    img,
    model,
    affinity_mode=affinity_mode,
    patch_shape=patch_shape,
)

# Generate segmentation via watershed and supervoxel agglomeration
segmentation = inference.affinities_to_segmentation(affinites)

# Save result as ZIP archive of SWC files
zip_path = "path-to-zip"
inference.segmentation_to_zipped_swcs(segmentation, zip_path)

Train

Here's an example of training a model,

from aind_exaspim_neuron_segmentation.machine_learning.data_handling import TrainDataset, ValidateDataset
from aind_exaspim_neuron_segmentation.machine_learning.train import Trainer


# Initializations
affinity_mode = True
patch_shape = (96, 96, 96)
output_path = "directory-path-to-save-ckpts"

# Paths
train_img_paths = "list-of-train-input-image-paths"
train_label_paths = "list-of-train-label-mask-paths"
val_img_paths = "list-of-validation-input-image-paths"
val_label_paths = "list-of-validation-label-mask-paths"

# Datasets
train_dataset = TrainDataset(
   train_img_paths,
   train_label_paths,
   affinity_mode=affinity_mode,
   patch_shape=patch_shape,
)
val_dataset = ValidateDataset(
   val_img_paths,
   val_label_paths,
   affinity_mode=affinity_mode,
   patch_shape=patch_shape
)

# Train
trainer = Trainer(output_dir)
trainer.run(train_dataset, val_dataset)

Contact Information

For any inquiries, feedback, or contributions, please do not hesitate to contact us. You can reach us via email at anna.grim@alleninstitute.org or connect on LinkedIn.

License

aind-exaspim-neuron-segmentation is licensed under the MIT License.