Add supervised from-scratch j0126 affinity checkpoint (MedNeXt-L, [48,96,96]) + model card

README.md

@@ -1,3 +1,88 @@
 ---
 license: mit
+tags:
+  - connectomics
+  - electron-microscopy
+  - 3d-segmentation
+  - affinity-prediction
+  - pytorch-connectomics
+  - mednext
 ---
+
+# j0126 affinity model (supervised, from scratch)
+
+3-channel voxel affinity prediction for the **j0126 zebra finch** EM volume, for use with
+[PyTorch Connectomics](https://github.com/PytorchConnectomics/pytorch_connectomics)
+(`tutorials/neuron_j0126`).
+
+| | |
+|---|---|
+| File | `affinity_scratch_48x96x96.ckpt` (247 MB) |
+| Architecture | MedNeXt-L, kernel 3 (61.8 M parameters) |
+| Input / output | 1 channel EM → 6 channel affinity (`banis+`; channels 0–2 are the nearest-neighbour affinity used downstream) |
+| Training patch | `[48, 96, 96]` (ZYX) — anisotropic, ≈ isotropic in nm for the 9×9×20 nm grid |
+| Training data | the 33 dense-GT j0126 cubes, heavy `aug_em_neuron` augmentation |
+| Schedule | 200 k steps from scratch (no pretrained init), AdamW, lr 1e-3, batch 8 × 4 GPUs |
+
+> **This model is not ground-truth-free.** It was trained on labelled j0126 tissue and is
+> the *supervised upper reference* for the zero-shot pipeline in the tutorial, not a
+> drop-in replacement for it. (The dense cubes come from a different alignment than the
+> 50 evaluation skeletons, so it is not voxel-level test contamination — but it is
+> supervision on the target domain.)
+
+## Inference window must be `[48, 96, 96]`
+
+MedNeXt normalizes with `GroupNorm(num_groups=C, num_channels=C)` — per-sample,
+per-channel instance norm with **no running statistics** — so normalization statistics are
+computed over the sliding window's spatial extent at every block and **the forward pass is
+window-size dependent**. Running this checkpoint at the tutorial's inherited
+`[144, 144, 144]` is a 3× Z / 1.5× XY mismatch that inverts the trained Z-thin anisotropy.
+
+`tutorials/neuron_j0126/1_affinity_supervised.yaml` already sets both
+`model.input_size` and `inference.window.window_size` to `[48, 96, 96]`. Both must be set
+in YAML, never on the CLI: the stage merge can clobber a pre-resolution CLI override.
+
+## Usage
+
+```bash
+huggingface-cli download pytc/j0126 affinity_scratch_48x96x96.ckpt --local-dir ckpt/
+
+python scripts/main.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml \
+  --mode test --checkpoint ckpt/affinity_scratch_48x96x96.ckpt
+```
+
+Whole volume (726 chunks of 1008³ with a 72-voxel halo), one GPU per shard:
+
+```bash
+python scripts/main.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml \
+  --mode test --checkpoint ckpt/affinity_scratch_48x96x96.ckpt \
+  --shard-id "$SLURM_ARRAY_TASK_ID" --num-shards 726
+```
+
+Output is one float16 CZYX HDF5 per chunk. Affinities are written with the `scale_sigmoid`
+activation, which the tutorial's ABISS watershed thresholds are calibrated against
+(`scale_sigmoid` 0.70 == plain sigmoid 0.9857) — do not substitute a plain sigmoid without
+recalibrating step 2.
+
+## Results
+
+Full-volume j0126, 50 evaluation skeletons, from `tutorials/neuron_j0126/README.md`. This
+checkpoint is the "scratch" affinity in that table:
+
+| Affinity | Decoding | Error correction | NERL mt=0 ↑ | NERL mt=5 ↑ | VOI ↓ |
+|---|---|---|---:|---:|---:|
+| FFN reference | — | — | **0.526** | 0.538 | **1.856** |
+| scratch | ABISS, exclusion mask | — | 0.268 | 0.470 | 2.584 |
+| scratch | + nucleus instance certificate | — | 0.287 | 0.482 | 2.562 |
+| scratch | + nucleus instance certificate | morphology-guided branch linking | 0.301 | **0.539** | 2.374 |
+| scratch | + nucleus instance certificate | + 3×3×3 inter-object erosion | **0.441** | 0.528 | 2.440 |
+
+Affinity quality alone is not what these numbers measure — see the tutorial for the
+decode and error-correction steps that produce them.
+
+## Checkpoint format
+
+A PyTorch Lightning checkpoint (`state_dict` keys are prefixed `model.model.`). Optimizer
+state has been stripped for download size; the weights are bit-identical to the reference
+run's `step=00200000.ckpt` (all 517 tensors, float32). It is loadable for inference/finetuning,
+not for resuming the original optimizer trajectory.

affinity_scratch_48x96x96.ckpt

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