Upload dbMiM neuron segmentation weights and model card

README.md

@@ -0,0 +1,313 @@
+# dbMiM Neuron Segmentation
+
+[![Chinese README](https://img.shields.io/badge/README-%E4%B8%AD%E6%96%87-blue)](README_zh.md)
+[![Hugging Face Weights](https://img.shields.io/badge/HuggingFace-weights-yellow)](https://huggingface.co/che111/dbmim-neuron-segmentation)
+
+This repository is the cleaned implementation used for our current dbMiM
+neuron-segmentation experiments on CREMI. The maintained path is:
+
+- self-supervised dbMiM / MAE-style pretraining on unlabeled EM volumes;
+- anisotropic 3D UNETR affinity finetuning on CREMI;
+- full-volume CREMI A/B/C evaluation with VOI and adapted Rand error (ARAND);
+- waterz-based post-processing with calibration and threshold sweeps.
+
+The old private cluster launchers, scratch reports, cached bytecode, legacy
+models, and historical dataloaders have been removed from Git. Large data,
+checkpoints, reports, and local experiment outputs are intentionally ignored.
+
+## Current Method
+
+The current best method is not the original minimal reproduction. It combines
+the following changes that were stable in ablations:
+
+1. **Anisotropic UNETR backbone**: `UNETRAnisotropicAffinityNet` with
+   `32x160x160` input crops, `patch_size=(4,16,16)`, transformer hidden states
+   used as UNETR skips, staged decoder upsampling, and a z-only anisotropic
+   transition before the final decoder block.
+2. **dbMiM pretraining on EM volumes**: a ViT/MAE encoder is pretrained with
+   masked reconstruction, membrane-aware weighting, and a lightweight structure
+   consistency loss. The pretrained encoder keys are then loaded into UNETR.
+3. **MSE + MAWS supervised finetuning**: CREMI labels are converted to z/y/x
+   nearest-neighbor affinities. Finetuning uses pure MSE with membrane-aware
+   spatial weighting (MAWS), channel weights `[1.35, 1.0, 1.0]`, and synchronized
+   image/label augmentations.
+4. **Official-style waterz evaluation**: full CREMI A/B/C labeled volumes are
+   evaluated with `ignore_label=0`, CREMI-style XY boundary ignore distance `1`,
+   z boundary ignore `0`, logit calibration biases, and a waterz threshold sweep.
+
+The most useful new finding is that **fixed mixed edge/random masking on full
+EM data (R33)** improves over scratch, old fullEM dbMiM, pure edge masking, and
+fullEM plain MAE. The best absolute VOI is still the smaller publicEM dbMiM
+model (R17), so the README reports both.
+
+## Model Zoo
+
+Weights are hosted at:
+
+**https://huggingface.co/che111/dbmim-neuron-segmentation**
+
+| Model | HF path | Intended use |
+|---|---|---|
+| PublicEM dbMiM R17 pretrain | `weights/publicem_dbmim_r17/pretrained_latest.pt` | ViT/dbMiM encoder checkpoint for UNETR initialization |
+| PublicEM dbMiM R17 finetune | `weights/publicem_dbmim_r17/finetuned_latest.pt` | Best current publicEM segmentation checkpoint |
+| FullEM mixed-mask dbMiM R33 pretrain | `weights/fullem_mixedmask_dbmim_r33/pretrained_latest.pt` | Recommended full-data dbMiM pretraining checkpoint |
+| FullEM mixed-mask dbMiM R33 finetune | `weights/fullem_mixedmask_dbmim_r33/finetuned_latest.pt` | Recommended full-data segmentation checkpoint |
+
+The pretraining checkpoints contain the masked-image-modeling encoder and
+decoder state. During finetuning we load only compatible encoder prefixes
+(`pos_embed`, `patch_embed`, `encoder_blocks`, `norm`) into the anisotropic
+UNETR. The finetuned checkpoints are full affinity segmentation models.
+
+## Data
+
+### Supervised CREMI Data
+
+Finetuning and evaluation use the public labeled CREMI 2016 training volumes:
+
+```text
+data/CREMI/sample_A_20160501.hdf
+data/CREMI/sample_B_20160501.hdf
+data/CREMI/sample_C_20160501.hdf
+```
+
+The raw key is `volumes/raw`; the instance-label key is
+`volumes/labels/neuron_ids`.
+
+### Pretraining Data
+
+Two unlabeled EM pretraining sets were used.
+
+| Name | Contents | Config examples |
+|---|---|---|
+| publicEM | CREMI raw + public ISBI 2012 + SNEMI3D raw volumes | `configs/pretrain_public_em_membrane_r16.yaml`, `configs/pretrain_public_em_plain_mae_r23.yaml` |
+| fullEM | CREMI raw + `cyd0806/EM_pretrain_data` groups: FAFB, FIB-25, Kasthuri, MitoEM, MB-MOC | `configs/pretrain_em_full_mixedmask_dbmim_r33.yaml`, `configs/pretrain_em_full_plain_mae_r23.yaml` |
+
+No hidden CREMI challenge labels are used in this repository.
+
+### Evaluation Split
+
+The reported numbers are **official-style validation on the public labeled
+CREMI A/B/C training volumes**, not challenge-server hidden-test results.
+
+The protocol is:
+
+- train supervised affinity models from random crops sampled from CREMI A/B/C;
+- run full-volume sliding-window inference on A, B, and C;
+- apply CREMI-style boundary ignore with `xy=1`, `z=0`;
+- sweep calibration biases and waterz thresholds;
+- report aggregate A/B/C `voi_sum` and `adapted_rand_error`.
+
+This split is small but matches the controlled ablation goal: isolate whether
+dbMiM pretraining improves the same anisotropic UNETR finetuning recipe over
+scratch and plain MAE controls.
+
+## Results
+
+Lower VOI and lower ARAND are better. `ARAND at best VOI` is the adapted Rand
+error at the threshold selected by lowest VOI. `Best ARAND` is selected
+independently and is included because VOI and ARAND can prefer different
+post-processing thresholds.
+
+### PublicEM Pretraining
+
+| Arm | VOI | ARAND at best VOI | Best ARAND | Conclusion |
+|---|---:|---:|---:|---|
+| R17 publicEM random-mask dbMiM | **1.002919** | **0.188832** | 0.188832 | Best publicEM VOI |
+| R23 publicEM random-mask plain MAE | 1.027073 | 0.192763 | 0.189247 | Matched MAE baseline |
+| R29 publicEM pure edge-mask dbMiM | 1.033564 | 0.186827 | **0.186827** | Best publicEM ARAND, worse VOI |
+| R32 publicEM fixed mixed-mask dbMiM | 1.046538 | 0.206256 | 0.193183 | Negative vs R17/R23 |
+| R34 publicEM adaptive mixed dbMiM | 1.067471 | 0.205437 | 0.200604 | Negative adaptive result |
+| R30 publicEM pure edge-mask plain MAE | 1.077594 | 0.203182 | 0.198562 | Edge-mask MAE control |
+| R17 scratch UNETR | 1.095164 | 0.213401 | 0.210442 | Scratch control |
+
+Key deltas:
+
+- R17 dbMiM beats matched publicEM plain MAE R23 by `-0.0242` VOI and about
+  `-0.0004` best ARAND.
+- R29 edge-mask dbMiM beats same-mask plain MAE R30 by `-0.0440` VOI and
+  `-0.0117` best ARAND, but its VOI is worse than R17/R23.
+
+### FullEM Pretraining
+
+| Arm | VOI | ARAND at best VOI | Best ARAND | Conclusion |
+|---|---:|---:|---:|---|
+| R33 fullEM fixed mixed-mask dbMiM | **1.039372** | **0.191216** | **0.190932** | Best fullEM result |
+| R31 fullEM pure edge-mask dbMiM | 1.055438 | 0.195125 | 0.195125 | Positive but weaker than R33 |
+| R20 fullEM old dbMiM | 1.085331 | 0.195722 | 0.195722 | Older fullEM baseline |
+| R35 fullEM adaptive mixed dbMiM | 1.089639 | 0.205551 | 0.205551 | Negative vs R33/R31/R20 |
+| R17 scratch UNETR | 1.095164 | 0.213401 | 0.210442 | Scratch control |
+| R23 fullEM plain MAE | 1.440684 | 0.281216 | 0.281216 | Negative fullEM MAE baseline |
+
+Key deltas:
+
+- R33 fullEM mixed-mask dbMiM beats fullEM plain MAE R23 by `-0.4013` VOI and
+  `-0.0903` best ARAND.
+- R33 beats scratch by `-0.0558` VOI and `-0.0195` best ARAND.
+- R33 beats old fullEM R20 by about `-0.0460` VOI.
+- R33 is still slightly worse than R17 publicEM by VOI (`1.039372` vs
+  `1.002919`), so the full-data recipe is the best fullEM result but not the
+  best global checkpoint yet.
+
+### Adaptive Masking
+
+R34/R35 tested an adaptive mixed masking policy that chooses mask ratio and
+edge fraction per crop. It did not improve downstream segmentation. After step
+40k, the policy collapsed to sampled mask ratio `0.75`; the mean learned edge
+fraction was `0.4456` for R34 and `0.3322` for R35. The current adaptive policy
+is therefore kept as a negative ablation rather than the recommended method.
+
+## Training Strategy
+
+### Pretraining
+
+Representative command:
+
+```bash
+python train_pretrain.py \
+  --config configs/pretrain_em_full_mixedmask_dbmim_r33.yaml
+```
+
+Main settings:
+
+| Setting | Value |
+|---|---|
+| Crop | `32x160x160` |
+| Patch size | `4x16x16` |
+| Encoder | ViT, `embed_dim=192`, `depth=6`, `heads=6` |
+| Mask ratio | `0.75` |
+| R33 mask strategy | `edge_random_mix`, `edge_mask_fraction=0.5`, `edge_mask_power=1.25` |
+| dbMiM losses | reconstruction + structure loss `0.2` + membrane weighting `1.35` |
+| Batch size | 2 per GPU |
+| Schedule | 160k optimizer steps, AdamW, lr `1.5e-4`, weight decay `0.05`, AMP |
+
+Plain MAE controls set `architecture: plain_mae`, `structure_weight: 0.0`, and
+`membrane_weight: 0.0` while keeping data, crop, model size, mask ratio, and
+schedule matched.
+
+### Finetuning
+
+Representative command:
+
+```bash
+python train_finetune.py \
+  --config configs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q.yaml
+```
+
+Main settings:
+
+| Setting | Value |
+|---|---|
+| Backbone | `unetr_aniso_em` |
+| Output | 3 affinity channels: z, y, x |
+| Crop | `32x160x160` |
+| Loss | MSE + MAWS, no BCE/Dice in the current winning recipe |
+| Label handling | synchronized image/label augmentation, 2D border widening radius 1 |
+| Batch size | 2 per GPU |
+| Schedule | 12k optimizer steps, lr `8e-5`, encoder lr `1e-5`, weight decay `0.01`, AMP |
+| Pretrained prefixes | `pos_embed`, `patch_embed`, `encoder_blocks`, `norm` |
+
+### Evaluation
+
+Representative full-volume command:
+
+```bash
+python scripts/evaluate_cremi_segmentation.py \
+  --config configs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q.yaml \
+  --checkpoint outputs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q/finetuned_latest.pt \
+  --data-dir data/CREMI \
+  --output-dir outputs/eval_cremi_r33_waterz_abc \
+  --crop-size 0 0 0 \
+  --stride 16 80 80 \
+  --backends waterz \
+  --thresholds 0.35 0.40 0.45 0.50 0.55 \
+  --calibration-biases -0.50 -1.00 -1.00 -0.25 -0.50 -0.50 0.0 0.0 0.0 \
+  --metric-backend skimage \
+  --ignore-label 0 \
+  --cremi-boundary-ignore-distance-xy 1 \
+  --cremi-boundary-ignore-distance-z 0 \
+  --max-samples 0 \
+  --device cuda
+```
+
+The evaluation writes:
+
+```text
+cremi_segmentation_records.json
+cremi_segmentation_metrics.csv
+cremi_segmentation_summary.json
+```
+
+Use `best_by_voi_sum` for headline VOI and inspect `best_by_adapted_rand` as a
+separate ARAND-selected operating point.
+
+## Quick Start
+
+Install the Python dependencies:
+
+```bash
+pip install -r requirements-dbMIM.txt
+```
+
+Run the synthetic smoke test:
+
+```bash
+bash scripts/run_smoke.sh
+```
+
+Compile the maintained entry points:
+
+```bash
+python -m py_compile \
+  dbmim/*.py \
+  train_pretrain.py \
+  train_finetune.py \
+  scripts/download_data.py \
+  scripts/inspect_hdf5.py \
+  scripts/prepare_public_em_pretrain_data.py \
+  scripts/prepare_em_pretrain_data.py \
+  scripts/evaluate_cremi_segmentation.py \
+  scripts/evaluate_cremi_blockwise_scale.py
+```
+
+Download weights from Hugging Face with `huggingface_hub`:
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+    repo_id="che111/dbmim-neuron-segmentation",
+    local_dir="outputs/hf_weights",
+    allow_patterns=["weights/**", "configs/**"],
+)
+```
+
+## Repository Layout
+
+```text
+dbmim/                         Core datasets, models, metrics, post-processing, utilities
+configs/                       Maintained smoke, recommended, and matched ablation configs
+scripts/download_data.py       CREMI download helper
+scripts/prepare_*_data.py      PublicEM / fullEM pretraining data preparation helpers
+scripts/evaluate_*.py          VOI/ARAND and blockwise-scale evaluation
+train_pretrain.py              dbMiM / MAE pretraining entry point
+train_finetune.py              affinity finetuning entry point
+requirements-dbMIM.txt         Python dependency list
+```
+
+## Citation
+
+```bibtex
+@inproceedings{chen2023self,
+  title={Self-supervised neuron segmentation with multi-agent reinforcement learning},
+  author={Chen, Yinda and Huang, Wei and Zhou, Shenglong and Chen, Qi and Xiong, Zhiwei},
+  booktitle={Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence},
+  pages={609--617},
+  year={2023}
+}
+```
+
+## Data and Credential Notes
+
+Use every external EM dataset under its original license and access policy.
+Do not commit downloaded datasets, generated checkpoints, TOS credentials,
+Hugging Face tokens, GitHub tokens, or cluster credentials.

README_zh.md

@@ -0,0 +1,293 @@
+# dbMiM 神经元分割
+
+[![English README](https://img.shields.io/badge/README-English-blue)](README.md)
+[![Hugging Face 权重](https://img.shields.io/badge/HuggingFace-weights-yellow)](https://huggingface.co/che111/dbmim-neuron-segmentation)
+
+这个仓库是当前 dbMiM 神经元分割实验的清理版代码，主要面向 CREMI
+复现、预训练收益验证、UNETR 结构改进和 VOI/ARAND 评测。当前维护的主线是：
+
+- 在未标注 EM 体数据上做 dbMiM / MAE 风格自监督预训练；
+- 在 CREMI 上用各向异性 3D UNETR 做 affinity 微调；
+- 在 CREMI A/B/C 全体积上用 VOI 和 adapted Rand error (ARAND) 评测；
+- 用 waterz 后处理，并做 calibration / threshold sweep。
+
+旧的私有集群提交器、历史实验报告、缓存字节码、旧 dataloader、旧模型和临时脚本已经从 Git 中删除。大数据、权重、报告和本地实验输出会被 `.gitignore` 忽略。
+
+## 当前方法
+
+当前最好方法不是最早的简单复现，而是经过消融后稳定有效的一条主线：
+
+1. **各向异性 UNETR backbone**：使用 `UNETRAnisotropicAffinityNet`，
+   输入 crop 为 `32x160x160`，`patch_size=(4,16,16)`，Transformer hidden
+   states 作为 UNETR skip，decoder 分阶段上采样，并在最后 decoder 前加入
+   z-only 的各向异性 transition。
+2. **EM 体数据上的 dbMiM 预训练**：先训练 ViT/MAE encoder，目标包括 masked
+   reconstruction、membrane-aware weighting 和轻量 structure consistency loss，
+   再把 encoder 权重加载到 UNETR。
+3. **MSE + MAWS 微调**：CREMI instance label 被转换成 z/y/x 三通道最近邻
+   affinity。当前获胜微调方案使用纯 MSE，加 membrane-aware spatial weighting
+   (MAWS)，通道权重为 `[1.35, 1.0, 1.0]`，并保证图像和 label 的几何增强同步。
+4. **official-style waterz 评测**：在 CREMI A/B/C 全体积上评测，使用
+   `ignore_label=0`，CREMI 风格 XY boundary ignore distance `1`，z 方向为 `0`，
+   并做 logit calibration bias 与 waterz threshold sweep。
+
+最新有价值的正结果是：**fullEM 数据上的固定 mixed edge/random masking (R33)**
+相比 scratch、旧 fullEM dbMiM、纯 edge mask 和 fullEM plain MAE 都有提升。
+但当前全局最低 VOI 仍然来自更小 publicEM 数据上的 R17，所以文档同时报告
+R17 和 R33。
+
+## 权重
+
+训练好的预训练权重和微调权重已经上传到：
+
+**https://huggingface.co/che111/dbmim-neuron-segmentation**
+
+| 模型 | HF 路径 | 用途 |
+|---|---|---|
+| PublicEM dbMiM R17 预训练 | `weights/publicem_dbmim_r17/pretrained_latest.pt` | ViT/dbMiM encoder 初始化权重 |
+| PublicEM dbMiM R17 微调 | `weights/publicem_dbmim_r17/finetuned_latest.pt` | 当前 publicEM 最好分割权重 |
+| FullEM mixed-mask dbMiM R33 预训练 | `weights/fullem_mixedmask_dbmim_r33/pretrained_latest.pt` | 推荐的 full-data dbMiM 预训练权重 |
+| FullEM mixed-mask dbMiM R33 微调 | `weights/fullem_mixedmask_dbmim_r33/finetuned_latest.pt` | 推荐的 full-data 分割权重 |
+
+预训练 checkpoint 包含 masked-image-modeling 的 encoder/decoder 状态。微调时只把兼容的 encoder 前缀加载到各向异性 UNETR：`pos_embed`、`patch_embed`、`encoder_blocks`、`norm`。微调 checkpoint 是完整 affinity 分割网络。
+
+## 数据
+
+### 有标注 CREMI 数据
+
+微调和评测使用公开的 CREMI 2016 training volumes：
+
+```text
+data/CREMI/sample_A_20160501.hdf
+data/CREMI/sample_B_20160501.hdf
+data/CREMI/sample_C_20160501.hdf
+```
+
+raw key 是 `volumes/raw`，instance label key 是
+`volumes/labels/neuron_ids`。
+
+### 预训练数据
+
+使用过两套未标注 EM 预训练数据。
+
+| 名称 | 内容 | 配置示例 |
+|---|---|---|
+| publicEM | CREMI raw + public ISBI 2012 + SNEMI3D raw volumes | `configs/pretrain_public_em_membrane_r16.yaml`, `configs/pretrain_public_em_plain_mae_r23.yaml` |
+| fullEM | CREMI raw + `cyd0806/EM_pretrain_data` 的 FAFB、FIB-25、Kasthuri、MitoEM、MB-MOC | `configs/pretrain_em_full_mixedmask_dbmim_r33.yaml`, `configs/pretrain_em_full_plain_mae_r23.yaml` |
+
+本仓库没有使用 CREMI challenge hidden test labels。
+
+### 训练和评测划分
+
+本文档中的结果是 **public labeled CREMI A/B/C training volumes 上的 official-style validation**，不是 challenge server hidden-test 结果。
+
+评测口径如下：
+
+- 监督微调时从 CREMI A/B/C 中随机采样 crop；
+- 评测时对 A/B/C 三个体数据做 full-volume sliding-window inference；
+- metric 计算时启用 CREMI-style boundary ignore：`xy=1`，`z=0`；
+- sweep calibration bias 和 waterz threshold；
+- 报告 A/B/C 聚合后的 `voi_sum` 和 `adapted_rand_error`。
+
+这个划分很小，但适合当前目标：在同一各向异性 UNETR、同一微调 recipe、同一后处理下，对比 dbMiM 预训练、scratch 和 plain MAE 控制组。
+
+## 结果
+
+VOI 和 ARAND 都是越低越好。`ARAND at best VOI` 是 VOI 最优阈值对应的
+ARAND；`Best ARAND` 是单独按 ARAND 选出的最优阈值，因为 VOI 和 ARAND 有时会偏好不同后处理点。
+
+### PublicEM 预训练
+
+| 实验 | VOI | ARAND at best VOI | Best ARAND | 结论 |
+|---|---:|---:|---:|---|
+| R17 publicEM random-mask dbMiM | **1.002919** | **0.188832** | 0.188832 | publicEM 最好 VOI |
+| R23 publicEM random-mask plain MAE | 1.027073 | 0.192763 | 0.189247 | matched MAE baseline |
+| R29 publicEM pure edge-mask dbMiM | 1.033564 | 0.186827 | **0.186827** | publicEM 最好 ARAND，但 VOI 较差 |
+| R32 publicEM fixed mixed-mask dbMiM | 1.046538 | 0.206256 | 0.193183 | 相比 R17/R23 为负 |
+| R34 publicEM adaptive mixed dbMiM | 1.067471 | 0.205437 | 0.200604 | adaptive 为负 |
+| R30 publicEM pure edge-mask plain MAE | 1.077594 | 0.203182 | 0.198562 | edge-mask MAE 控制组 |
+| R17 scratch UNETR | 1.095164 | 0.213401 | 0.210442 | scratch 控制组 |
+
+关键差值：
+
+- R17 dbMiM 相比 matched publicEM plain MAE R23：VOI 降低 `0.0242`，best ARAND 约降低 `0.0004`。
+- R29 edge-mask dbMiM 相比同 mask 的 plain MAE R30：VOI 降低 `0.0440`，best ARAND 降低 `0.0117`，但 VOI 不如 R17/R23。
+
+### FullEM 预训练
+
+| 实验 | VOI | ARAND at best VOI | Best ARAND | 结论 |
+|---|---:|---:|---:|---|
+| R33 fullEM fixed mixed-mask dbMiM | **1.039372** | **0.191216** | **0.190932** | 最好 fullEM 结果 |
+| R31 fullEM pure edge-mask dbMiM | 1.055438 | 0.195125 | 0.195125 | 正收益，但弱于 R33 |
+| R20 fullEM old dbMiM | 1.085331 | 0.195722 | 0.195722 | 旧 fullEM baseline |
+| R35 fullEM adaptive mixed dbMiM | 1.089639 | 0.205551 | 0.205551 | 弱于 R33/R31/R20 |
+| R17 scratch UNETR | 1.095164 | 0.213401 | 0.210442 | scratch 控制组 |
+| R23 fullEM plain MAE | 1.440684 | 0.281216 | 0.281216 | fullEM MAE 明显为负 |
+
+关键差值：
+
+- R33 fullEM mixed-mask dbMiM 相比 fullEM plain MAE R23：VOI 降低 `0.4013`，best ARAND 降低 `0.0903`。
+- R33 相比 scratch：VOI 降低 `0.0558`，best ARAND 降低 `0.0195`。
+- R33 相比旧 fullEM R20：VOI 约降低 `0.0460`。
+- R33 仍然略差于 publicEM R17 的最好 VOI (`1.039372` vs `1.002919`)，所以 fullEM recipe 是当前最好 fullEM 方案，但还不是全局最好 checkpoint。
+
+### Adaptive Masking
+
+R34/R35 测试了每个 crop 自适应选择 mask ratio 和 edge fraction 的 mixed masking
+policy。这个方向目前没有带来提升。40k step 之后，policy 基本收敛到
+`sampled_mask_ratio=0.75`；R34 平均 `edge_fraction=0.4456`，R35 平均
+`edge_fraction=0.3322`。因此当前 adaptive policy 作为负消融保留，不作为推荐方法。
+
+## 训练策略
+
+### 预训练
+
+代表性命令：
+
+```bash
+python train_pretrain.py \
+  --config configs/pretrain_em_full_mixedmask_dbmim_r33.yaml
+```
+
+主要设置：
+
+| 设置 | 值 |
+|---|---|
+| Crop | `32x160x160` |
+| Patch size | `4x16x16` |
+| Encoder | ViT, `embed_dim=192`, `depth=6`, `heads=6` |
+| Mask ratio | `0.75` |
+| R33 mask strategy | `edge_random_mix`, `edge_mask_fraction=0.5`, `edge_mask_power=1.25` |
+| dbMiM loss | reconstruction + structure loss `0.2` + membrane weighting `1.35` |
+| Batch size | 每张 GPU 2 |
+| Schedule | 160k optimizer steps, AdamW, lr `1.5e-4`, weight decay `0.05`, AMP |
+
+plain MAE 控制组使用 `architecture: plain_mae`，`structure_weight: 0.0`，
+`membrane_weight: 0.0`，其它数据、crop、模型大小、mask ratio 和 schedule 保持匹配。
+
+### 微调
+
+代表性命令：
+
+```bash
+python train_finetune.py \
+  --config configs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q.yaml
+```
+
+主要设置：
+
+| 设置 | 值 |
+|---|---|
+| Backbone | `unetr_aniso_em` |
+| 输出 | 3 个 affinity channel：z、y、x |
+| Crop | `32x160x160` |
+| Loss | MSE + MAWS，当前获胜 recipe 不用 BCE/Dice |
+| Label 处理 | 图像/label 几何增强同步，2D border widening radius 1 |
+| Batch size | 每张 GPU 2 |
+| Schedule | 12k optimizer steps, lr `8e-5`, encoder lr `1e-5`, weight decay `0.01`, AMP |
+| 预训练加载前缀 | `pos_embed`, `patch_embed`, `encoder_blocks`, `norm` |
+
+### 评测
+
+代表性 full-volume 命令：
+
+```bash
+python scripts/evaluate_cremi_segmentation.py \
+  --config configs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q.yaml \
+  --checkpoint outputs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q/finetuned_latest.pt \
+  --data-dir data/CREMI \
+  --output-dir outputs/eval_cremi_r33_waterz_abc \
+  --crop-size 0 0 0 \
+  --stride 16 80 80 \
+  --backends waterz \
+  --thresholds 0.35 0.40 0.45 0.50 0.55 \
+  --calibration-biases -0.50 -1.00 -1.00 -0.25 -0.50 -0.50 0.0 0.0 0.0 \
+  --metric-backend skimage \
+  --ignore-label 0 \
+  --cremi-boundary-ignore-distance-xy 1 \
+  --cremi-boundary-ignore-distance-z 0 \
+  --max-samples 0 \
+  --device cuda
+```
+
+评测会输出：
+
+```text
+cremi_segmentation_records.json
+cremi_segmentation_metrics.csv
+cremi_segmentation_summary.json
+```
+
+主表使用 `best_by_voi_sum` 作为 VOI 汇报点，同时检查 `best_by_adapted_rand`
+以避��只看单一阈值。
+
+## 快速开始
+
+安装依赖：
+
+```bash
+pip install -r requirements-dbMIM.txt
+```
+
+运行 synthetic smoke test：
+
+```bash
+bash scripts/run_smoke.sh
+```
+
+编译当前维护入口：
+
+```bash
+python -m py_compile \
+  dbmim/*.py \
+  train_pretrain.py \
+  train_finetune.py \
+  scripts/download_data.py \
+  scripts/inspect_hdf5.py \
+  scripts/prepare_public_em_pretrain_data.py \
+  scripts/prepare_em_pretrain_data.py \
+  scripts/evaluate_cremi_segmentation.py \
+  scripts/evaluate_cremi_blockwise_scale.py
+```
+
+用 `huggingface_hub` 下载权重：
+
+```python
+from huggingface_hub import snapshot_download
+
+snapshot_download(
+    repo_id="che111/dbmim-neuron-segmentation",
+    local_dir="outputs/hf_weights",
+    allow_patterns=["weights/**", "configs/**"],
+)
+```
+
+## 仓库结构
+
+```text
+dbmim/                         核心 dataset、model、metric、postprocess、utils
+configs/                       当前保留的 smoke、推荐方法和 matched ablation 配置
+scripts/download_data.py       CREMI 下载辅助脚本
+scripts/prepare_*_data.py      publicEM / fullEM 预训练数据准备脚本
+scripts/evaluate_*.py          VOI/ARAND 和 blockwise-scale 评测脚本
+train_pretrain.py              dbMiM / MAE 预训练入口
+train_finetune.py              affinity 微调入口
+requirements-dbMIM.txt         Python 依赖
+```
+
+## 引用
+
+```bibtex
+@inproceedings{chen2023self,
+  title={Self-supervised neuron segmentation with multi-agent reinforcement learning},
+  author={Chen, Yinda and Huang, Wei and Zhou, Shenglong and Chen, Qi and Xiong, Zhiwei},
+  booktitle={Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence},
+  pages={609--617},
+  year={2023}
+}
+```
+
+## 数据和密钥说明
+
+外部 EM 数据请遵守其原始 license 和访问规则。不要把下载的数据、生成的
+checkpoint、TOS 凭证、Hugging Face token、GitHub token 或集群凭证提交到仓库。

checksums.sha256

@@ -0,0 +1,4 @@
+8b94dcbb4a6bed3590de145ddaa2604a1888b2dc11fea7d7e4ce860cd4c2cdd6  weights/publicem_dbmim_r17/pretrained_latest.pt
+da2775fe7a0dbeaaa3ae68fabeb4ef786405dd1ff90a37eea30f157912adfd1c  weights/publicem_dbmim_r17/finetuned_latest.pt
+a86917528db0cf0106886e91c1f8564b8657f8c771209e6fda93f4374a10f364  weights/fullem_mixedmask_dbmim_r33/pretrained_latest.pt
+fbdd97de7c9e4c27f21ffe74f12664829360e1d6b2395665fe1ac85458d11ecd  weights/fullem_mixedmask_dbmim_r33/finetuned_latest.pt

configs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q.yaml

@@ -0,0 +1,100 @@
+seed: 295
+device: cuda
+output_dir: /volume/med-train/users/dchen02/code/dbMiM/outputs/finetune_cremi_real_unetr_aniso_em_mse_maws_fullem_mixedmask_r33q
+pretrained: /volume/med-train/users/dchen02/code/dbMiM/outputs/pretrain_em_full_mixedmask_dbmim_r33/pretrained_latest.pt
+pretrained_include_prefixes:
+- pos_embed
+- patch_embed
+- encoder_blocks
+- norm
+data:
+  synthetic: false
+  image_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  label_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  image_keys:
+  - volumes/raw
+  - raw
+  - main
+  label_keys:
+  - volumes/labels/neuron_ids
+  - labels
+  - label
+  - gt
+  volume_size:
+  - 32
+  - 160
+  - 160
+  length_multiplier: 2048
+  augment: true
+  widen_border: true
+  widen_border_radius: 1
+  augment_rotate_xy: true
+  augment_gamma: true
+  augment_gamma_range:
+  - 0.7
+  - 1.5
+  augment_noise_std: 0.035
+model:
+  architecture: unetr_aniso_em
+  in_channels: 1
+  out_channels: 3
+  volume_size:
+  - 32
+  - 160
+  - 160
+  patch_size:
+  - 4
+  - 16
+  - 16
+  embed_dim: 192
+  depth: 6
+  num_heads: 6
+  feature_size: 32
+  dropout: 0.05
+  em_refine_depth: 2
+  channel_bias_init:
+  - -0.2
+  - 0.0
+  - 0.0
+train:
+  batch_size: 2
+  epochs: 200
+  max_steps: 12000
+  num_workers: 8
+  lr: 8.0e-05
+  weight_decay: 0.01
+  amp: true
+  log_every: 20
+  eval_every: 0
+  eval_max_batches: 0
+  save_every: 0
+  save_steps: 1000
+  val_fraction: 0.0
+  clip_grad: 1.0
+  replicate_affinity_boundary: true
+  loss:
+    loss_type: mse
+    bce_weight: 1.0
+    dice_weight: 0.0
+    boundary_dice_weight: 0.0
+    bcar_weight: 0.0
+    bcar_calibration_weight: 0.0
+    membrane_weight: 0.75
+    membrane_axis_weights:
+    - 0.25
+    - 1.0
+    - 1.0
+    membrane_clip: 4.0
+    membrane_normalize: true
+    channel_weights:
+    - 1.35
+    - 1.0
+    - 1.0
+  encoder_lr: 1.0e-05
+  encoder_param_prefixes:
+  - pos_embed
+  - patch_embed
+  - encoder_blocks
+  - norm

configs/finetune_cremi_real_unetr_aniso_em_mse_maws_publicem_r17q.yaml

@@ -0,0 +1,95 @@
+seed: 191
+device: cuda
+output_dir: /volume/med-train/users/dchen02/code/dbMiM/outputs/finetune_cremi_real_unetr_aniso_em_mse_maws_publicem_r17q
+pretrained: /volume/med-train/users/dchen02/code/dbMiM/outputs/pretrain_public_em_membrane_dbmim_r16/pretrained_latest.pt
+data:
+  synthetic: false
+  image_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  label_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  image_keys:
+  - volumes/raw
+  - raw
+  - main
+  label_keys:
+  - volumes/labels/neuron_ids
+  - labels
+  - label
+  - gt
+  volume_size:
+  - 32
+  - 160
+  - 160
+  length_multiplier: 2048
+  augment: true
+  widen_border: true
+  widen_border_radius: 1
+  augment_rotate_xy: true
+  augment_gamma: true
+  augment_gamma_range:
+  - 0.7
+  - 1.5
+  augment_noise_std: 0.035
+model:
+  architecture: unetr_aniso_em
+  in_channels: 1
+  out_channels: 3
+  volume_size:
+  - 32
+  - 160
+  - 160
+  patch_size:
+  - 4
+  - 16
+  - 16
+  embed_dim: 192
+  depth: 6
+  num_heads: 6
+  feature_size: 32
+  dropout: 0.05
+  em_refine_depth: 2
+  channel_bias_init:
+  - -0.2
+  - 0.0
+  - 0.0
+train:
+  batch_size: 2
+  epochs: 200
+  max_steps: 12000
+  num_workers: 8
+  lr: 8.0e-05
+  weight_decay: 0.01
+  amp: true
+  log_every: 20
+  eval_every: 0
+  eval_max_batches: 0
+  save_every: 0
+  save_steps: 1000
+  val_fraction: 0.0
+  clip_grad: 1.0
+  replicate_affinity_boundary: true
+  loss:
+    loss_type: mse
+    bce_weight: 1.0
+    dice_weight: 0.0
+    boundary_dice_weight: 0.0
+    bcar_weight: 0.0
+    bcar_calibration_weight: 0.0
+    membrane_weight: 0.75
+    membrane_axis_weights:
+    - 0.25
+    - 1.0
+    - 1.0
+    membrane_clip: 4.0
+    membrane_normalize: true
+    channel_weights:
+    - 1.35
+    - 1.0
+    - 1.0
+  encoder_lr: 1.0e-05
+  encoder_param_prefixes:
+  - pos_embed
+  - patch_embed
+  - encoder_blocks
+  - norm

configs/pretrain_em_full_mixedmask_dbmim_r33.yaml

@@ -0,0 +1,70 @@
+seed: 295
+device: cuda
+output_dir: /volume/med-train/users/dchen02/code/dbMiM/outputs/pretrain_em_full_mixedmask_dbmim_r33
+data:
+  synthetic: false
+  train_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  - /volume/med-train/users/dchen02/code/dbMiM/data/EM_pretrain_data/all
+  image_keys:
+  - volumes/raw
+  - raw
+  - main
+  - data
+  volume_size:
+  - 32
+  - 160
+  - 160
+  length_multiplier: 8192
+  augment: true
+  augment_rotate_xy: true
+  augment_gamma: true
+  augment_gamma_range:
+  - 0.7
+  - 1.5
+  augment_noise_std: 0.025
+model:
+  architecture: dbmim
+  in_channels: 1
+  volume_size:
+  - 32
+  - 160
+  - 160
+  patch_size:
+  - 4
+  - 16
+  - 16
+  embed_dim: 192
+  depth: 6
+  num_heads: 6
+  decoder_dim: 192
+  mask_ratio: 0.75
+  mask_strategy: edge_random_mix
+  edge_mask_fraction: 0.5
+  edge_mask_power: 1.25
+  edge_mask_noise: 0.05
+  structure_weight: 0.2
+  structure_axis_weights:
+  - 0.5
+  - 1.0
+  - 1.0
+  membrane_weight: 1.35
+  membrane_axis_weights:
+  - 0.25
+  - 1.0
+  - 1.0
+  membrane_clip: 5.0
+decision:
+  enabled: false
+train:
+  batch_size: 2
+  epochs: 200
+  max_steps: 160000
+  num_workers: 8
+  lr: 0.00015
+  weight_decay: 0.05
+  amp: true
+  log_every: 20
+  save_every: 5
+  save_steps: 2000
+  clip_grad: 1.0

configs/pretrain_public_em_membrane_r16.yaml

@@ -0,0 +1,72 @@
+seed: 176
+device: cuda
+output_dir: /volume/med-train/users/dchen02/code/dbMiM/outputs/pretrain_public_em_membrane_dbmim_r16
+data:
+  synthetic: false
+  train_paths:
+  - /volume/med-train/users/dchen02/code/dbMiM/data/CREMI
+  - /volume/med-train/users/dchen02/code/dbMiM/data/EM_pretrain_data/public_em
+  image_keys:
+  - volumes/raw
+  - raw
+  - main
+  - data
+  volume_size:
+  - 32
+  - 160
+  - 160
+  length_multiplier: 8192
+  augment: true
+  augment_rotate_xy: true
+  augment_gamma: true
+  augment_gamma_range:
+  - 0.7
+  - 1.5
+  augment_noise_std: 0.025
+model:
+  in_channels: 1
+  volume_size:
+  - 32
+  - 160
+  - 160
+  patch_size:
+  - 4
+  - 16
+  - 16
+  embed_dim: 192
+  depth: 6
+  num_heads: 6
+  decoder_dim: 192
+  mask_ratio: 0.75
+  structure_weight: 0.2
+  structure_axis_weights:
+  - 0.5
+  - 1.0
+  - 1.0
+  membrane_weight: 1.35
+  membrane_axis_weights:
+  - 0.25
+  - 1.0
+  - 1.0
+  membrane_clip: 5.0
+decision:
+  enabled: true
+  hidden_dim: 256
+  target_mask_ratio: 0.75
+  min_mask_ratio: 0.4
+  max_mask_ratio: 0.9
+  lr: 0.0005
+  policy_weight: 0.05
+  freeze_after_steps: 40000
+train:
+  batch_size: 2
+  epochs: 200
+  max_steps: 160000
+  num_workers: 8
+  lr: 0.00015
+  weight_decay: 0.05
+  amp: true
+  log_every: 20
+  save_every: 5
+  save_steps: 2000
+  clip_grad: 1.0

weights/fullem_mixedmask_dbmim_r33/finetuned_latest.pt

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

weights/fullem_mixedmask_dbmim_r33/pretrained_latest.pt

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

weights/publicem_dbmim_r17/finetuned_latest.pt

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

weights/publicem_dbmim_r17/pretrained_latest.pt

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