pytc/tutorial

PyTorch Connectomics pretrained models

Pretrained checkpoints for 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 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
neuron_nisb/	neuron instance segmentation	BANIS+ (MedNeXt-L + cc3d)	NISB base	NERL 0.601	neuron_nisb
mito_lucchi++/	mitochondria semantic segmentation	MedNeXt-S	Lucchi++	Jaccard 0.935	mito_lucchi++

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

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):

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:

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.