---
license: mit
library_name: pytorch-connectomics
pipeline_tag: image-segmentation
tags:
  - connectomics
  - electron-microscopy
  - em
  - neuron-segmentation
  - instance-segmentation
  - 3d
  - affinity
  - mednext
  - biology
  - neuroscience
  - pytorch
datasets:
  - yanchaoz/EMNeuron
---

# PyTorch Connectomics pretrained models

Pretrained checkpoints for [PyTorch Connectomics](https://github.com/PytorchConnectomics/pytorch_connectomics)
(PyTC) tutorials: 3D electron-microscopy (EM) segmentation of neurons and mitochondria. Each
folder holds one model; its PyTC tutorial has the config, the commands and the full results.

| folder | task | model | data | result | tutorial |
|---|---|---|---|---|---|
| [`neuron_segneuron/`](neuron_segneuron) | **neuron instance segmentation, any EM volume** (zero-shot or fine-tune) | BANIS+ (MedNeXt-L affinities + ABISS watershed) | EMNeuron, 10 labeled EM sources | mean VI 0.869 on SegNeuron's 8 evaluation volumes | [`neuron_segneuron`](https://github.com/PytorchConnectomics/pytorch_connectomics/tree/master/tutorials/neuron_segneuron) |
| [`neuron_nisb/`](neuron_nisb) | neuron instance segmentation | BANIS+ (MedNeXt-L + cc3d) | NISB base | NERL 0.601 | [`neuron_nisb`](https://github.com/PytorchConnectomics/pytorch_connectomics/tree/master/tutorials/neuron_nisb) |
| [`mito_lucchi++/`](mito_lucchi%2B%2B) | mitochondria semantic segmentation | MedNeXt-S | Lucchi++ | Jaccard 0.935 | [`mito_lucchi++`](https://github.com/PytorchConnectomics/pytorch_connectomics/tree/master/tutorials/mito_lucchi%2B%2B) |

## Segment neurons in your EM volume

`neuron_segneuron/banis_emneuron_200k.ckpt` is the general-purpose neuron model: trained on
EMNeuron, the EM database of [SegNeuron](https://github.com/yanchaoz/SegNeuron): 10 labeled
sources (AxonEM-H, AxonEM-M, CREMI, FIB-25, Hemibrain, J0126, Kasthuri, three Mira volumes),
6-9 nm in xy and 8-50 nm in z. Start here for a new EM volume; fine-tune if you have labels.

```bash
git clone https://github.com/PytorchConnectomics/pytorch_connectomics && cd pytorch_connectomics
# install: see QUICKSTART.md; ABISS (github.com/PytorchConnectomics/ABISS) decodes affinities
# input: a directory of <name>_img.h5 (one 3D uint8 dataset, ZYX)
ABISS_HOME=/path/to/abiss tutorials/neuron_segneuron/predict.sh /path/to/volumes /path/to/output
# -> /path/to/output/<name>_seg.h5 (neuron ids) and <name>_aff.h5 (affinities)
```

`predict.sh` downloads the checkpoint on first use. `TTA=1` adds 16-view test-time
augmentation; `MERGE=0.7` (default 0.6) splits more if neurons are merged.

Fine-tune on labeled volumes (`<name>_img.h5` + `<name>_seg.h5` instance labels):

```bash
tutorials/neuron_segneuron/finetune.sh /path/to/train /path/to/val /path/to/run
ABISS_HOME=/path/to/abiss tutorials/neuron_segneuron/predict.sh <printed checkpoint> /path/to/volumes /path/to/output
```

Download only the checkpoint:

```bash
hf download pytc/tutorial neuron_segneuron/banis_emneuron_200k.ckpt --local-dir checkpoints
```

## Citation

If you use these models, cite PyTorch Connectomics (arXiv:2112.05754); for the EMNeuron
model also cite SegNeuron (Zhang et al., MICCAI 2024), whose data it was trained on.