Add neuron_segneuron: BANIS+ trained on EMNeuron (200k steps, EMA); model card

Mean VI 0.870 / ARE 0.104 on SegNeuron 8 evaluation volumes (TTA16, ABISS merge 0.6 chosen on held-out training chunks). Tutorial: PytorchConnectomics/pytorch_connectomics tutorials/neuron_segneuron.

README.md

@@ -1,3 +1,67 @@
----
-license: mit
----
+---
+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.870 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.

neuron_segneuron/banis+.yaml

@@ -0,0 +1,156 @@
+# BANIS+ on EMNeuron, SegNeuron's train/val split, with the MedNeXt training
+# speed knobs.
+#
+#   data:   python scripts/prepare_emneuron.py --src /path/to/EMNeuron --out /path/to/emneuron_pytc
+#   train:  python scripts/main.py --config tutorials/neuron_segneuron/banis+.yaml
+#   infer:  python scripts/main.py --config tutorials/neuron_segneuron/banis+.yaml --mode test \
+#               --checkpoint <run>/checkpoints/last.ckpt
+#   score:  python scripts/evaluate_emneuron.py --pred-dir <run>/results --gt /path/to/EMNeuron/valid
+#
+# Recipe = ../banis+.yaml (6-channel affinity, MedNeXt-L/k3, per-channel
+# class-balanced BCE, EMA 0.999, erosion 2, AdamW + cosine, 200k steps) with
+# these changes:
+#
+# AFFINITY LAYOUT -- ABISS's, not BANIS's: channels x, y, z (axes 2, 1, 0 of
+# the ZYX arrays) and each edge (i-1, i) stored at voxel i (`deepem` mode,
+# ABISS "destination" storage). The predictions go to decode_abiss as they
+# are; the BANIS layout (z, y, x at the source voxel) needs a channel reversal
+# and a one-voxel shift first.
+#
+# PATCH -- EMNeuron mixes 6-9 nm xy with 8-50 nm z, and a third of its volumes
+# hold only 31-50 sections, so the 128^3 cube becomes [32, 128, 128] ZYX
+# (SegNeuron trains on 20 x 128 x 128). MedNeXt-L is stride-16: 32/16 = 2.
+# Long-range offsets are 10 voxels in x/y and 4 sections in z (BANIS uses 10 on
+# every axis of a 9x9x20 nm volume; 10 sections would span a third of the
+# patch). Destination storage needs the voxel at i - offset, so target_context
+# reads 4/10/10 voxels on the leading side.
+#
+# AUGMENTATION -- aug_banis permutes z with y/x and rotates in 3D, which is
+# wrong for anisotropic sections. aug_em_neuron_fast keeps flips, xy
+# quarter-turns, contrast/brightness/noise and one section defect per sample;
+# motion blur rejoins the defect set, and an xy-only scale jitter (+-20%,
+# independent per axis) stands in for SegNeuron's anisotropic scaling. No 3D
+# elastic: its isotropic displacement would move 30-50 nm sections by many slices.
+#
+# SPEED -- the MedNeXt training knobs: torch.compile with the shift-sum
+# depthwise conv, fused AdamW with merged parameter groups (bf16, so no fp16
+# GradScaler blocks fused clipping), cudnn.benchmark, the numba affine
+# resampler, and a frozen MedNeXt dummy tensor so DDP runs with
+# find_unused_parameters=false.
+
+_base_:
+- ../banis+.yaml
+- dataset.yaml
+
+experiment_name: emneuron_banis+
+description: >-
+  BANIS+ (6-channel affinity, MedNeXt-L/k3, class-balanced BCE, EMA, erosion 2)
+  on SegNeuron's EMNeuron split: ABISS-layout affinities decoded with ABISS,
+  [32,128,128] patches, anisotropic EM augmentation, and the MedNeXt speed
+  knobs (compile + shift-sum depthwise, fused AdamW, bf16).
+
+default:
+  system:
+    num_workers: 8
+  model:
+    input_size: [32, 128, 128]
+    output_size: [32, 128, 128]
+    compile: true
+    mednext_depthwise_impl: shift_sum
+    mednext_freeze_dummy_tensor: true
+  data:
+    dataloader:
+      # 308 small chunks: preload them all (~17 GB per DDP rank); banis.yaml's
+      # `lazy` profile targets the large NISB zarr.
+      profile: cached
+      use_lazy_zarr: false
+      use_cache: false
+      use_preloaded_cache_train: true
+      use_preloaded_cache_val: true
+      patch_size: [32, 128, 128]
+      target_context: [4, 10, 10, 0, 0, 0]
+      batch_size: 4
+    label_transform:
+      targets:
+      - name: affinity
+        kwargs:
+          offsets: ["0-0-1", "0-1-0", "1-0-0", "0-0-10", "0-10-0", "4-0-0"]
+          affinity_mode: deepem
+    augmentation:
+      profile: aug_em_neuron_fast
+      motion_blur:
+        enabled: true
+        prob: 1.0
+        sections: [0, 5]
+        kernel_size: 11
+        sigma_range: [1.0, 3.0]
+        full_section_prob: 0.5
+        partial_ratio_range: [0.25, 0.75]
+      affine:
+        enabled: true
+        prob: 0.5
+        rotate_range: [0.0, 0.0, 0.0]
+        scale_range: [0.0, 0.2, 0.2]
+        shear_range: [0.0, 0.0, 0.0]
+        fast_resample: true
+  inference:
+    # Predict with the EMA weights the recipe trains and validates.
+    checkpoint_weights: ema
+    model:
+      # Short-range x, y, z: already ABISS's channel order.
+      select_channel: [0, 1, 2]
+    window:
+      # Infer at the trained window: MedNeXt's per-sample GroupNorm makes the
+      # forward pass depend on the window extent.
+      window_size: [32, 128, 128]
+      sw_batch_size: 16
+  decoding:
+    steps:
+    # ABISS ws: watershed, size-dependent single linkage, dust removal. The
+    # runner's default edge storage is "destination", which this model writes.
+    # Thresholds selected on the 10 tune/ chunks (TTA16, lowest mean VI) by
+    # scripts/evaluate_emneuron.py, which also fills voxels ABISS leaves
+    # unlabeled with the nearest segment before scoring.
+    - template: decoding_abiss
+      kwargs:
+        abiss_home: /path/to/abiss
+        cli_args:
+          ws_high_threshold: 0.9
+          ws_low_threshold: 0.05
+          ws_size_threshold: 1000000000
+          ws_dust_threshold: 0
+          ws_merge_function: max
+          ws_merge_threshold: 0.6
+  evaluation:
+    # VI / adapted Rand are scored by scripts/evaluate_emneuron.py with the
+    # SegNeuron evaluator's exact calls; the volumes have no NERL skeletons.
+    enabled: false
+
+test:
+  # Destination storage leaves the first voxel of every axis without an inside
+  # edge, so the test pipeline crops it from deepem predictions. One voxel of
+  # reflect context per side, minus the trailing voxel here, keeps the saved
+  # affinity on the input grid: ABISS reads it with no shift or pad.
+  data:
+    data_transform:
+      pad_size: [1, 1, 1]
+  inference:
+    model:
+      crop_pad: [0, 1, 0, 1, 0, 1]
+
+train:
+  optimization:
+    precision: "bf16-mixed"
+    benchmark: true
+    optimizer:
+      fused: true
+      merge_param_groups: true
+    ddp_find_unused_parameters: false
+    val_check_interval: 10000
+    val_steps_per_epoch: 50
+  monitor:
+    checkpoint:
+      # Step checkpoints every 25k plus last.ckpt; the reported model is the
+      # final step, so the val-loss top-1 copy is kept for diagnostics only.
+      save_top_k: 1
+      save_every_n_steps: 25000

neuron_segneuron/banis_emneuron_200k.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b763de15418698c08d057da95060a536d12c8a136c8a99e0fb8faf65d6ab914
+size 494572982