SegForge

English | 中文

Extension framework for nnU-Net v2: custom trainers, network factories, and discovery CLIs.

Recommended usage: install SegForge as a library, point nnUNet_extTrainer at the installed built-in trainers, and keep experiment-specific -tr classes in your own project folder (do not fork SegForge per dataset).

Requirements

• Python ≥ 3.10
• nnU-Net v2 (nnunetv2≥2.8.1) and a working nnU-Net data layout (nnUNet_raw / nnUNet_preprocessed / nnUNet_results)

Recommended setup

1. Install as a library

# From PyPI (recommended for users)
pip install segforge
pip install "segforge[neuron]"     # NeuronSWC / npu
pip install "segforge[mamba]"      # U-Mamba / LightM / LKM / nnMamba
pip install "segforge[segmamba]"   # SegMamba BiMamba-v3
pip install "segforge[opaque]"     # Nuitka; hide project network .py as .so
# pip install "segforge[full]"     # mamba + segmamba + scipy / scikit-image

# From a local checkout (SegForge development)
pip install -e ".[dev]"

Extras detail: docs/networks.md.

2. Wire nnU-Net to SegForge trainers (+ optional project trainers)

# Built-in trainers from the installed package
export SEGFORGE_TR="$(python -c 'import segforge.training.nnUNetTrainer as t; print(t.__path__[0])')"

# Recommended for experiments: also append your project trainer dir
export nnUNet_extTrainer="${SEGFORGE_TR}:/path/to/my_project/trainers"
# Built-ins only:
# export nnUNet_extTrainer="${SEGFORGE_TR}"

(: on Linux/macOS, ; on Windows.)

3. List and train

segforge_list_trainers -s external    # SegForge built-ins + project trainers
nnUNetv2_train DATASET_ID 3d_fullres 0 -tr YourProjectTrainer

Binary {0, 255} export (segmentation; remap at write time — no second pass):

# same flags as nnUNetv2_predict
segforge_predict_binary255 -d DATASET_ID -c 3d_fullres -f 0 -i INPUT -o OUTPUT -tr YourTrainer

Distance / soft-map regression predict (default float32 continuous; do not use stock nnUNetv2_predict — it truncates to uint8):

# continuous soft map (~[0,1] with default sigmoid)
segforge_predict_regression -d DATASET_ID -c 3d_fullres -f 0 -i INPUT -o OUTPUT \
  -p nnUNetPlans_reg -tr nnUNetDistanceTrainer

# also dump continuous .npz
segforge_predict_regression ... --save_probabilities

# thresholded binary {0,255} or {0,1} (can combine with --save_probabilities)
segforge_predict_regression ... --threshold 0.5 --binary-value 255
segforge_predict_regression ... --threshold 0.5 --binary-value 1

Project trainers import SegForge bases / mixins / existing trainers; they should not live inside the SegForge source tree unless you are contributing upstream. Details: docs/customizing.md.

Quick start

segforge_list_trainers -s external
segforge_list_mixins -g optim
segforge_list_networks -b monai

nnUNetv2_train DATASET_ID 3d_fullres 0 -tr UNetTrainer

segforge_list_trainers / segforge_list_mixins show a NOTES column from class docstrings (summary + optional paper: / dataset:; use -v for full text). segforge_list_networks -v prints module/builder details.

Networks

Builders, -tr names, CUDA / 2D–3D constraints, and SegMamba BiMamba notes: docs/networks.md.

Highlight (recent): MedNeXt-v2, SegMamba ([segmamba]), SegFormer3D, D-LKA (2D), UNETR++, CMUNeXt (2D).

Trainer layout

Naming (recommended): Model[_Patch][_Opt][_LR][_Schedule] — say what changed, not the paper dataset. Put provenance in docstring tags (paper: / dataset:).

Hyperparams: use FixedHyperMixin (or any Adam* mixin that inherits it) for fixed_initial_lr / fixed_weight_decay / fixed_num_epochs.

Compose order: mixins on the left, architecture trainer on the right:

from segforge.training.bases import (
    AdamWCosineNoClipMixin,
    FixedPatchMixin,
    NoDeepSupervisionTrainer,
    SoftclDiceCELossMixin,
    SystemBenchmarkMixin,
)

# Single-output net + forced isotropic patch + AdamW cosine (FixedHyper included)
class FooTrainer_128_AdamW1en3_Cosine(
    AdamWCosineNoClipMixin, FixedPatchMixin, NoDeepSupervisionTrainer
):
    """Foo; patch 128; AdamW lr=1e-3 + cosine.

    paper: …
    """

    fixed_patch_size = 128
    fixed_initial_lr = 1e-3

# Hang a loss recipe on an existing model trainer
from segforge.training.nnUNetTrainer.CoTrTrainer import CoTrTrainer_128
from segforge.training.nnUNetTrainer.MedNeXtTrainer import MedNeXtTrainer_S_kernel3

class CoTr_clDice(SoftclDiceCELossMixin, CoTrTrainer_128):
    pass

# Efficiency probe (not a stock -tr; compose then train briefly)
class MedNeXtBench(SystemBenchmarkMixin, MedNeXtTrainer_S_kernel3):
    benchmark_patch_size = 128

• Forces patch: name includes size and uses FixedPatchMixin (e.g. CoTrTrainer_128, VNetTrainer_96).
• Plans patch: no size in the class name (e.g. VNetTrainer, UNETRTrainer).

Dataset conversion

Challenge → nnU-Net raw helpers live under segforge/dataset_conversion/ (not on the training import path). Set nnUNet_raw, then:

python -m segforge.dataset_conversion.Parse2022 -i /path/to/PARSE2022 -d 501
python -m segforge.dataset_conversion.LiTS2017 -i /path/to/LiTS2017 -d 502
python -m segforge.dataset_conversion.DRIVE -i /path/to/DRIVE -d 503
python -m segforge.dataset_conversion.AIIB2023 -i /path/to/AIIB2023 -d 504
python -m segforge.dataset_conversion.BTCV -i /path/to/RawData -d 505
# Neuron: image + SWC → mask or dist
# pip install "segforge[neuron]"   # installs neuron-process-utils (npu) from PyPI
python -m segforge.dataset_conversion.NeuronSWC --image /data/img --swc /data/swc -d 601 --label-type mask
python -m segforge.dataset_conversion.NeuronSWC --image /data/img --swc /data/swc -d 602 --label-type dist -t soft_centerline
# -d is required (nnU-Net folder: DatasetXXX_<Name>)

Docs

• Changelog
• Networks & optional extras
• Custom networks, losses & trainer composition
• Distance-map regression
• Hide unpublished network source (.so)
• 中文：网络与可选依赖 · 自定义模型、损失与 Trainer 组合 · 距离图回归 · 隐藏未发表网络源码

中文版：README.zh-CN.md

License

GPL-3.0-only. See LICENSE.

SegForge builds on nnU-Net and third-party architectures/losses; cite the original papers when you use those components. Some vendored files retain their original upstream licenses (e.g. Apache-2.0 CoTr Deformable DETR snippets, MIT CMUNeXt); the overall SegForge distribution is GPL-3.0. Details: docs/networks.md.