Reorganise into per-dataset folders: move the j0126 affinity model to j0126/

README.md

@@ -9,113 +9,20 @@ tags:
   - mednext
 ---
 
-# j0126 affinity model (supervised, from scratch)
+# pytc/vEM — volume EM models
 
-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`).
+Checkpoints for [PyTorch Connectomics](https://github.com/PytorchConnectomics/pytorch_connectomics),
+one folder per dataset. Each folder has its own README with the training recipe, the
+inference constraints that matter, and the measured results.
 
-| | |
-|---|---|
-| Files | `affinity_scratch_48x96x96.ckpt` (247 MB) · `affinity_scratch_48x96x96.onnx` (247 MB, opset 17) |
-| 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
+| Folder | Dataset | Task | Files |
+|---|---|---|---|
+| [`j0126/`](j0126) | j0126 zebra finch, 9 × 9 × 20 nm | 3-channel voxel affinity | `affinity_scratch_48x96x96.ckpt` (247 MB) · `.onnx` (247 MB, opset 17) |
 
 ```bash
-hf 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
+hf download pytc/vEM j0126/affinity_scratch_48x96x96.ckpt --local-dir 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.
-
-## ONNX
-
-`affinity_scratch_48x96x96.onnx` is the same weights exported for runtimes without PyTorch.
-
-```python
-import numpy as np, onnxruntime as ort
-
-session = ort.InferenceSession("affinity_scratch_48x96x96.onnx",
-                               providers=["CPUExecutionProvider"])
-image = np.random.randn(1, 1, 48, 96, 96).astype(np.float32)   # (B, C, Z, Y, X)
-logits = session.run(["logits"], {"image": image})[0]          # (B, 6, 48, 96, 96)
-affinity = 1.0 / (1.0 + np.exp(-0.2 * logits))                 # scale_sigmoid, temperature 0.2
-```
-
-Three things the graph does **not** do for you:
-
-- **The activation.** The output is raw logits. The pipeline applies `scale_sigmoid` with
-  BANIS' temperature 0.2 — `sigmoid(0.2 * x)` — downstream, and step 2's watershed thresholds
-  are calibrated against that.
-- **Channel selection.** Six channels are emitted; the decode consumes channels 0–2.
-- **Sliding-window assembly.** One window in, one window out. Tiling, overlap and blending
-  are the caller's job.
-
-The spatial dims are fixed at `[48, 96, 96]` and only the batch axis is dynamic — see the
-window section above for why a dynamic spatial axis would be a trap here. Deep supervision is
-switched off in the export, so the graph has a single output.
-
-Verified against PyTorch on identical input: max |diff| 2.2e-04 on logits spanning
-[1.03, 21.17] (relative 1.0e-05), which is **4.6e-06** after `scale_sigmoid` — float32
-accumulation noise, not a behavioural difference. Regenerate with
-`python scripts/export_onnx.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml
---checkpoint <ckpt> --output model.onnx --check`.
-
-## 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.
+> Renamed from `pytc/j0126` on 2026-09-08 and reorganised into per-dataset folders. The old
+> repo URL redirects, but a link that pinned `affinity_scratch_48x96x96.ckpt` at the repo
+> root needs the `j0126/` prefix added.

j0126/README.md

@@ -0,0 +1,119 @@
+# 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`).
+
+| | |
+|---|---|
+| Files | `affinity_scratch_48x96x96.ckpt` (247 MB) · `affinity_scratch_48x96x96.onnx` (247 MB, opset 17) |
+| 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.
+
+The tutorial sets this in one place: `inference.window_size: [48, 96, 96]` in
+`tutorials/neuron_j0126/params.yaml`, which `2_infer.yaml` reads. Set it in YAML, never on
+the CLI: the stage merge can clobber a pre-resolution CLI override.
+
+## Usage
+
+The tutorial downloads this checkpoint for you: leave `train.enabled: false` in
+`tutorials/neuron_j0126/params.yaml` and run the whole pipeline with one command.
+
+```bash
+python scripts/run_j0126.py
+```
+
+To use it directly instead:
+
+```bash
+hf download pytc/vEM j0126/affinity_scratch_48x96x96.ckpt --local-dir ckpt/
+
+python scripts/main.py --config tutorials/neuron_j0126/2_infer.yaml \
+  --mode test --checkpoint ckpt/j0126/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/2_infer.yaml \
+  --mode test --checkpoint ckpt/j0126/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 3.
+
+## 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.
+
+## ONNX
+
+`affinity_scratch_48x96x96.onnx` is the same weights exported for runtimes without PyTorch.
+
+```python
+import numpy as np, onnxruntime as ort
+
+session = ort.InferenceSession("affinity_scratch_48x96x96.onnx",
+                               providers=["CPUExecutionProvider"])
+image = np.random.randn(1, 1, 48, 96, 96).astype(np.float32)   # (B, C, Z, Y, X)
+logits = session.run(["logits"], {"image": image})[0]          # (B, 6, 48, 96, 96)
+affinity = 1.0 / (1.0 + np.exp(-0.2 * logits))                 # scale_sigmoid, temperature 0.2
+```
+
+Three things the graph does **not** do for you:
+
+- **The activation.** The output is raw logits. The pipeline applies `scale_sigmoid` with
+  BANIS' temperature 0.2 — `sigmoid(0.2 * x)` — downstream, and step 2's watershed thresholds
+  are calibrated against that.
+- **Channel selection.** Six channels are emitted; the decode consumes channels 0–2.
+- **Sliding-window assembly.** One window in, one window out. Tiling, overlap and blending
+  are the caller's job.
+
+The spatial dims are fixed at `[48, 96, 96]` and only the batch axis is dynamic — see the
+window section above for why a dynamic spatial axis would be a trap here. Deep supervision is
+switched off in the export, so the graph has a single output.
+
+Verified against PyTorch on identical input: max |diff| 2.2e-04 on logits spanning
+[1.03, 21.17] (relative 1.0e-05), which is **4.6e-06** after `scale_sigmoid` — float32
+accumulation noise, not a behavioural difference. Regenerate with
+`python scripts/export_onnx.py --config tutorials/neuron_j0126/1_train.yaml
+--checkpoint <ckpt> --output model.onnx --check`.
+
+## 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.