juntaoJianggavin/APRIL-MedSeg

A Modern Modular 2D Medical Image Segmentation Toolbox

52

stars

45

commits

Python

primary language

Aug 14, 2026

updated

Browse cluster: Medical Image Segmentation with Foundation Models

README

APRIL-MedSeg Logo

License: Apache 2.0 arXiv

Juntao Jiang, Jinsheng Bai, Linxuan Fan, Yali Bi, Jiangning Zhang, Yong Liu

中文文档

130 networks · 177 encoders · 45 decoders · 81 losses · 25 skip connections · 17 bottlenecks · 6 training paradigms · 24 augmentations · 917 YAML configs · switch anything with one line of YAML


📰 Updates

  • 2026.07.30Bug fix: semi_train.py selected the best checkpoint with the EMA teacher but saved only the student weights, so best_model.pth could not reproduce the reported validation Dice for teacher-evaluated methods. The evaluated model is now stored in model_state_dict, and the student is kept in student_state_dict for resuming.
  • 2026.07.26 — Fixed RWKV-UNet WKV CUDA kernel loading and aligned GLSP with the official pretrained weights.
  • 2026.07.24Bug fix: Fixed multi-class mask loading across all datasets; pixel values are now correctly mapped to contiguous class indices.
  • 2026.07.20Bug fix: Fixed a critical data split issue in GenericDataset that caused incorrect train/val/test sample counts. All training/evaluation scripts updated. New models: Added support for SAM3 (Perception Encoder backbone), MedSAM2 (2D, Hiera backbone), and MedSAM3 (2D, PE backbone).
  • 2026.06.15 — The project has been officially renamed to APRIL-MedSeg, named after the APRIL Lab at Zhejiang University (led by Prof. Yong Liu).
  • 2026.06.11UltimateMedSeg is officially released!

📑 Table of Contents


📦 Installation

Requirements

  • Python >= 3.10
  • PyTorch >= 2.0
  • NVIDIA GPU with CUDA 12.4+

Basic Installation

git clone https://github.com/juntaoJianggavin/APRIL-MedSeg.git
cd APRIL-MedSeg

# Install dependencies
pip install -r requirements.txt

# Install in dev mode
pip install -e .

Optional Dependencies

# Foundation models
pip install transformers safetensors

# MLLM inference pipeline
pip install groundingdino-py
pip install git+https://github.com/facebookresearch/segment-anything.git

# ONNX export & verification
pip install onnx onnxruntime

# Lion optimizer
pip install lion-pytorch

# Mamba / SSM encoders (only required for Mamba-family networks; needs CUDA toolchain to build)
pip install causal-conv1d
pip install mamba-ssm

Pretrained Weights & Transfer Learning

Three levels of weight loading, from lightest to heaviest:

MethodYAML KeyScopeUse Case
Encoder pretrainedencoder.pretrained: trueBackbone onlyImageNet / domain-specific backbone weights
Manual pretrained pathencoder.pretrained_path: /path/to/weights.pthBackbone onlyOffline / custom backbone checkpoint
Transfer learningmodel.transfer_learning_path: /path/to/full_model.pthEntire networkLoad full model (encoder+decoder+bottleneck+head) from a previous training run
model:
  # Full-model transfer learning (loads AFTER encoder pretrained, takes priority)
  transfer_learning_path: null  # or /path/to/checkpoint.pth

  encoder:
    name: timm_resnet50
    pretrained: true              # auto-download ImageNet backbone weights
    pretrained_path: null         # or /path/to/backbone.pth (manual override)

Note: Models requiring specific pretrained weights (43 architectures in REQUIRES_PRETRAINED) will display a 10-second warning when pretrained: false. All auto-downloadable weights support manual path override via pretrained_path or transfer_learning_path.

Pretrained Weight Sources (4 Categories)

CategoryMechanismModels
A. WEIGHT_REGISTRY auto-downloadensure_weight() with GitHub/GCS/HF sourcesswinunet, h2former, hiformer, transunet, vm_unet, rwkv_unet (B/S/T), cswin_unet, da_transunet, mamba_unet, fcbformer, transnuseg
B. timm / torchvision runtimepretrained: true triggers built-in downloadsegformer_b0–b5, esfpnet, cascade, emcad, polyp_pvt, fatnet, transfuse, mist, hsnet, ssformer, ldnet, dconnnet, cfanet, lv_unet, nulite, polyper
C. SAM familypretrained: true auto-downloads ViT/SAM weightssam_b, sam_l, mobile_sam, sam2, sam_med2d, samed, sammed2d_wrapper, samus, auto_sam, lite_medsam, medical_sam_adapter
D. Foundation & MLLM encodersHuggingFace Hub / open_clip / transformers auto-download39 foundation encoders (dinov2, dino, clip_vit, sam_vit, dinov3, phikon, uni, plip, musk, phikon_v2, keep, raddino, omnirad, biovil, chexzero, retfound, retfound_dinov2, flair, ophmae, panderm, dermclip, monet_derm, endo_vit, endo_fm, surgical_sam, biomedclip, medclip, medsiglip, usfmae, ultrafedfm, samus) + 8 MLLM vision towers (qwen3_vl, qwen25_vl, llava_med, medgemma, healthgpt, huatuogpt, hulumed, lingshu)

Automatic Weight Download

# List all registered weights and cache status
python -m medseg.utils.weight_downloader list

# Download a specific weight (auto-retries all sources)
python -m medseg.utils.weight_downloader download medsam_vit_b

# Check which auto-downloadable weights are present
python -m medseg.utils.weight_downloader check

When all auto-download sources fail, the error message includes:

  • Manual download URL
  • Exact cache path to place the file at
  • Alternative: set pretrained_path in YAML to point to a local file

timm encoder weights are downloaded automatically via timm's built-in mechanism.


🚀 Quick Start

1. Standard Supervised Training

# ResNet50 + UNet decoder
python train.py --config configs/architectures/networks/general/aau_net.yaml \
    --output_dir output/aau_net

# With AMP mixed precision
python train.py --config configs/architectures/networks/general/transunet.yaml \
    --output_dir output/transunet --amp

# Multi-GPU DDP training
torchrun --nproc_per_node=4 train.py \
    --config configs/architectures/networks/general/swinunet.yaml \
    --output_dir output/swinunet --amp

2. Inference / Testing

# Single model evaluation
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth

# Save prediction results
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth --save_pred --output_dir test_output/

# Multi-checkpoint ensemble (logit averaging)
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint ckpt_a.pth ckpt_b.pth ckpt_c.pth \
    --ensemble-weights 0.5 0.3 0.2 \
    --ensemble-average logit

# Test-Time Augmentation (TTA)
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --tta \
    --tta-augs identity rot90 rot180 rot270 hflip vflip \
    --tta-merge mean

# TTA + Ensemble combined
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint ckpt_a.pth ckpt_b.pth \
    --ensemble-average logit \
    --tta --tta-merge mean

3. Semi-Supervised Training

# Mean Teacher
python semi_train.py --config configs/training_paradigms/semi_supervision/mean_teacher.yaml \
    --output_dir output/semi_mt

# CPS (Cross Pseudo Supervision)
python semi_train.py --config configs/training_paradigms/semi_supervision/cps.yaml \
    --output_dir output/semi_cps

4. Domain Adaptation Training

# AdvEnt
python train_domain_adaptation.py \
    --config configs/training_paradigms/domain_adaptation/advent.yaml \
    --output_dir output/da_advent

# TENT (Test-Time Adaptation)
python train_domain_adaptation.py \
    --config configs/training_paradigms/domain_adaptation/tent.yaml \
    --output_dir output/da_tent

5. Knowledge Distillation Training

python train_distillation.py \
    --teacher_config configs/training_paradigms/distillation/teacher_large.yaml \
    --student_config configs/training_paradigms/distillation/student_small.yaml \
    --distillation_type logit \
    --temperature 4.0 \
    --alpha 0.5 \
    --output_dir output/kd_logit

6. Weakly Supervised Training

# Box-supervised
python train_weakly_supervised.py \
    --config configs/training_paradigms/weak_supervision/box_supervised.yaml \
    --supervision_type box \
    --output_dir output/weak_box

# CAM-based
python train_weakly_supervised.py \
    --config configs/training_paradigms/weak_supervision/cam.yaml \
    --supervision_type cam \
    --output_dir output/weak_cam

7. Text-Guided Training

# Train
python train_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_clip.yaml \
    --output_dir output/text_cris

# Test (auto-detects: trainable model vs inference pipeline)
python test_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_clip.yaml \
    --checkpoint output/text_cris/best_model.pth

# Test inference-only pipeline (no checkpoint needed)
python test_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_grounding_dino_sam2.yaml

8. Model Profiling

# FLOPs / Params / FPS
python profile_model.py --config configs/architectures/networks/general/transunet.yaml

9. ONNX Export

python scripts/export_onnx.py \
    --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --output model.onnx --verify

10. Prediction Visualization

python scripts/visualize.py \
    --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --input ./data/test/images/ \
    --output vis_output/

11. Python API

from medseg.utils.config import load_config
from medseg.model_builder import build_model

cfg = load_config("configs/architectures/networks/general/transunet.yaml")
model = build_model(cfg)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable params: {trainable / 1e6:.2f}M")

📚 Tutorial

A step-by-step tutorial series covering deep learning medical image segmentation from fundamentals to advanced topics:

ChapterTitleKey Topics
01Introduction to Medical Image SegmentationConcepts, clinical significance, metrics, method evolution
02U-Net in DetailArchitecture, skip connections, U-Net family variants
03Data and PreprocessingFormats, split strategies, augmentation pipeline
04Training and EvaluationLoss functions, optimizers, AMP/DDP, evaluation
05Encoder Deep DiveCNN / Transformer / Mamba / RWKV comparison, timm wrapper
06Decoders and Skip ConnectionsCASCADE / EMCAD / Attention Gate, skip taxonomy
07Foundation ModelsDPT head, 9 medical modalities, fine-tuning strategies
08Advanced Training ParadigmsOverview: semi-supervised, domain adaptation, distillation, weakly supervised
08aSemi-Supervised SegmentationMean Teacher, CPS, UniMatch, FixMatch, consistency regularization
08bDomain AdaptationAdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo
08cKnowledge DistillationVanillaKD, DKD, CWD, MGD, DIST, ReviewKD, SimKD
08dWeakly Supervised SegmentationCAM, SEAM, PuzzleCAM, Box, Point, Scribble supervision
08eText-Guided SegmentationCRIS, BiomedParse, LViT, CLIP-based, MLLM pipeline
09Deployment and InferenceONNX export, TTA, ensemble, MLLM pipeline

Full tutorial index


🏗️ Project Structure

segmentation_tool/
├── medseg/                                      # Core framework
│   ├── models/                                  # Model components
│   │   ├── encoders/                            #   177 encoders (92 native + 85 timm presets + 1000+ via timm_ prefix)
│   │   │   ├── cnn/              (13 modules)   #     CNN: basic, DCSAU, CFA, MedNeXt, MEW, R2U, AttUNet, LV, MALU, EGE, ConvNeXt, EfficientNetV2, HRNet
│   │   │   ├── transformer/      (18 modules)   #     Transformer: TransUNet, SwinUNet, MISSFormer, DAEFormer, HiFormer, PVTv2, MaxViT, ViT-Pyramid, ...
│   │   │   ├── mamba/            (10 modules)   #     Mamba/SSM: VMUNet, UMamba, LKM, LoG-VMamba, UltraLight-VM, VMKLA, ...
│   │   │   ├── rwkv/             (5 modules)    #     RWKV: RWKV-UNet, U-RWKV (MICCAI), U-RWKV (TIP), MD-RWKV, RIR-Zigzag
│   │   │   ├── linear_attn/      (5 modules)    #     Linear attention: RetNet, Linformer, Performer, TTT, xLSTM
│   │   │   ├── kan_mlp/          (4 modules)    #     KAN/MLP: UKAN, Rolling-UNet, UNeXt, WA-UKAN
│   │   │   ├── foundation/       (39 modules)   #     Foundation models (DPT head)
│   │   │   │   ├── general/      (5)            #       DINOv2, DINOv3, DINO, CLIP-ViT, SAM-ViT
│   │   │   │   ├── pathology/    (6)            #       Phikon, Phikon-v2, UNI, PLIP, MUSK, KEEP
│   │   │   │   ├── radiology/    (4)            #       Rad-DINO, OmniRad, BioViL, CheXZero
│   │   │   │   ├── ophthalmology/(4)            #       RETFound-DINOv2, FLAIR, OphMAE, RETFound
│   │   │   │   ├── dermatology/  (3)            #       PanDerm, DermCLIP, MonetDerm
│   │   │   │   ├── general_medical/(3)          #       BiomedCLIP, MedCLIP, MedSigLIP
│   │   │   │   ├── mllm_vision/  (8)            #       Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, ...
│   │   │   │   ├── endoscopy/    (3)            #       EndoViT, Endo-FM, Surgical-SAM
│   │   │   │   └── ultrasound/   (3)            #       UltraFedFM, USF-MAE, SAMUS
│   │   │   └── wrapper/          (1 module)     #     timm dynamic wrapper (85 pre-registered + 1000+ via timm_ prefix)
│   │   ├── decoders/                            #   45 decoders
│   │   │   ├── basic/            (4 registered) #     Basic upsampling: deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep
│   │   │   ├── dense/            (2 registered) #     Dense connections: UNet++, UNet3+
│   │   │   ├── cascade/          (10 registered)#     CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet
│   │   │   ├── attention/        (6 registered) #     Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet
│   │   │   ├── transformer/      (5 registered) #     DAEFormer, MISSFormer, MTUNet, nnFormer, SwinUNet
│   │   │   ├── mlp/              (2 registered) #     SegFormer MLP, MLP Decoder
│   │   │   ├── specific/         (13 registered)#     CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet
│   │   │   ├── pyramid/          (2 registered) #     UPerNet, DeepLabV3 (ASPP)
│   │   │   └── mamba/            (1 registered) #     VM-UNet
│   │   ├── bottlenecks/          (17 modules)   #   17 bottlenecks: none, basic, ASPP, DenseASPP, PPM, Transformer, SE, CBAM, ...
│   │   ├── skip_connections/                    #   25 skip connections
│   │   │   ├── basic/            (3 modules)    #     Basic: concat, dense, add
│   │   │   ├── attention/        (10 modules)   #     Attention: AG, CAB, SAB, SCSE, CBAM, Gating, GRU, GAB, SC-Att, TA-MoSC
│   │   │   ├── transformer/      (5 modules)    #     Transformer: CrossAttn, TransFusion, AggAttn, MISSFormer, UCTrans
│   │   │   ├── mamba/            (1 module)     #     Mamba: SK-VM++
│   │   │   └── fusion/           (6 modules)    #     CNN fusion: BiFusion, Deformable, MultiScale, FeatureRefine, CCM, SDI
│   │   ├── networks/                            #   130 complete architectures (variants merged)
│   │   │   ├── cnn/              (36 registered)#     CNN: UNet, UNet3+, UNet++, AttUNet, nnUNet, MedNeXt, MEW-UNet, ...
│   │   │   ├── transformer/      (36 registered)#     Transformer: SegFormer, TransUNet, SwinUNet, DAEFormer, PolypPVT, CASCADE, ...
│   │   │   ├── mamba/            (24 registered)#     Mamba: VMUNet, U-Mamba, SwinUMamba, SkinMamba, DermoMamba, SerpMamba, ...
│   │   │   ├── sam/              (10 registered)#     SAM family: MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, ...
│   │   │   ├── rwkv/             (5 registered) #     RWKV: U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV, RIR-Zigzag
│   │   │   ├── kan_mlp/          (4 registered) #     KAN/MLP: RollingUNet, UNeXt, UKAN, WA-UKAN
│   │   │   └── linear_attn/      (3 registered) #     Linear attention: TTT-UNet, U-VixLSTM, xLSTM-UNet
│   │   └── text_unet/            (13 modules)   #   Text-guided (12 models): CRIS, BiomedParse, LanGuideMedSeg, LViT, TGANet, TPRO, ...
│   ├── training/                                # Training paradigms
│   │   ├── semi/                 (22 modules)   #   20 semi-supervised methods + 2 utils (base, utils)
│   │   │                                        #     MeanTeacher, CPS, UniMatch, FixMatch, UA-MT, CorrMatch, AllSpark, ...
│   │   ├── domain_adaptation/    (18 modules)   #   18 domain adaptation: AdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo, ...
│   │   ├── distillation/         (28 modules)   #   27 distillation: VanillaKD, DKD, MGD, DIST, CWD, ReviewKD, SimKD, NORM, ...
│   │   └── weakly_supervised/    (21 modules)   #   20 weakly supervised methods (CAM, SEAM, PuzzleCAM, TreeEnergy, ...)
│   ├── inference/                               # Inference
│   │   ├── ensemble.py                          #   Ensemble inference (multi-model voting)
│   │   ├── tta.py                               #   Test-time augmentation
│   │   └── mllm/                 (16 modules)   #   MLLM pipeline: 9 detector × 4 segmenter = 36 combinations
│   │       │                                    #     Detector: GroundingDINO, Qwen2/2.5/3-VL, InternVL, LLaVA, MiniCPM-V, Phi3-V, CogVLM
│   │       │                                    #     Segmenter: SAM2, MedSAM, SAM-Med2D, LiteMedSAM
│   │       └── medisee/          (3 modules)    #     MediSee: LLM reasoning segmenter
│   ├── losses/                   (15 modules)   # 81 losses
│   │                                            #   Supervised: CE, Dice, Focal, Tversky, Lovász, Boundary, Hausdorff, ...
│   │                                            #   Distillation: VanillaKD, DKD, CWD, MGD, DIST, AT, RKD, ...
│   │                                            #   Domain adaptation: AdvEnt, DANN, FDA, MIC, TENT, ...
│   │                                            #   Weakly supervised: Box, CAM, Point, Scribble, TreeEnergy, SEAM, ...
│   ├── datasets/                 (10 modules)   # Data loading: Synapse, ACDC, Generic, QaTa-COV19, MosMedData+, 24 augmentations
│   │   ├── advanced_aug.py                      #   24 advanced augmentations (YAML configurable)
│   │   └── transforms.py                        #   Basic transforms (Resize, ToTensor, Normalize)
│   ├── utils/                    (11 modules)   # Utilities
│   │   ├── amp_ddp.py                           #   AMP mixed precision + DDP distributed + DataParallel
│   │   ├── logger.py                            #   TensorBoard / WandB unified logging
│   │   ├── config.py                            #   Config inheritance (_base_ field support)
│   │   ├── warmup.py                            #   Warmup scheduler + Lion/AdamW/SGD optimizers
│   │   ├── augmentation.py                      #   Augmentation builder (basic/albumentations/pipeline)
│   │   ├── reproducibility.py                   #   Reproducibility (global seed + cuDNN deterministic)
│   │   ├── weight_downloader.py                 #   Automatic weight download + manual URL hints
│   │   ├── metrics.py                           #   Evaluation metrics: Dice, IoU, HD95, NSD
│   │   ├── hf_hub.py                            #   HuggingFace Hub model/dataset download
│   │   ├── timm_compat.py                       #   timm version compatibility utilities
│   │   └── timm_pretrained.py                   #   timm pretrained weight management
│   ├── text_guided.py                           # Text-guided segmentation (CRIS, BiomedParse, LanGuideMedSeg, ...)
│   ├── model_builder.py                         # YAML → model auto-assembler
│   └── registry.py                              # 6 registries: ENCODER / DECODER / SKIP / BOTTLENECK / LOSS / AUGMENTATION
├── data/                                        # Dataset root (user datasets go here)
│   ├── YourDataset/                             #   Your custom dataset
│   ├── source/                                  #   Domain adaptation source
│   ├── target/                                  #   Domain adaptation target
│   ├── target_val/                              #   Domain adaptation validation
│   └── test_dummy/                              #   Dummy test data
├── figs/                                        # Figures & logos
│   └── logo.png                                 #   Project logo
├── configs/                      (917 yamls)    # YAML configs
│   ├── architectures/            (783 yamls)    #   Network architecture configs
│   │   ├── networks/             (302 yamls)    #     Complete networks (130 arch across general/acdc/synapse)
│   │   ├── combinations/         (169 yamls)    #     Encoder+decoder free combinations
│   │   ├── decoder_study/        (133 yamls)    #     Decoder ablation (3 enc × 44 dec + 1)
│   │   ├── skip_study/           (75 yamls)     #     Skip ablation (3 enc × 25 skip)
│   │   ├── bottleneck_study/     (51 yamls)     #     Bottleneck ablation (3 enc × 17 bn)
│   │   └── foundation/           (53 yamls)     #     Foundation models (9 modalities × 39 encoders)
│   ├── training_paradigms/       (104 yamls)    #   Training paradigm configs
│   │   ├── semi_supervision/     (20 yamls)     #     Semi-supervised (20 methods)
│   │   ├── domain_adaptation/    (18 yamls)     #     Domain adaptation (18 methods)
│   │   ├── distillation/         (29 yamls)     #     Distillation (27 methods)
│   │   ├── text_guided/          (17 yamls)     #     Text-guided (12 models + pipeline)
│   │   └── weak_supervision/     (20 yamls)     #     Weakly supervised (20 methods)
│   └── intro_to_datasets/        (27 yamls)     #   27 dataset introductions + example configs
├── scripts/                                     # Utility + experiment scripts
│   ├── experiments/              (14 scripts)   #   Experiment bash scripts
│   │   ├── run_sota_benchmark.sh                #     SOTA architecture comparison (11 models × 7 datasets)
│   │   ├── run_decoder_study.sh                 #     Decoder ablation (3 enc × 15 classic dec)
│   │   ├── run_bottleneck_study.sh              #     Bottleneck ablation (3 enc × 9 bn)
│   │   ├── run_skip_study.sh                    #     Skip ablation (3 enc × 12 skip)
│   │   ├── run_polyp_benchmark.sh               #     Polyp-specific models (16 models × 2 datasets)
│   │   ├── run_skin_benchmark.sh                #     Skin-specific models (16 models × 2 datasets + PH2 external)
│   │   ├── run_retinal_benchmark.sh             #     Retinal-specific models (7 models × 3 datasets)
│   │   ├── run_ultrasound_benchmark.sh          #     Ultrasound-specific models (8 models × BUSI)
│   │   ├── run_pathology_benchmark.sh           #     Pathology-specific models (5 models × GlaS)
│   │   ├── run_lightweight_skin.sh              #     Lightweight skin segmentation (8 models)
│   │   ├── run_semi_study.sh                    #     Semi-supervised paradigm comparison (6 methods)
│   │   ├── run_da_study.sh                      #     Domain adaptation paradigm comparison (8 methods)
│   │   ├── run_kd_study.sh                      #     Knowledge distillation comparison (7 methods)
│   │   └── run_weak_study.sh                    #     Weakly supervised paradigm comparison (6 methods)
│   ├── check_config_paths.py                    #   Check config path references across docs/scripts
│   ├── download_hf_dataset.py                   #   Download HuggingFace datasets
│   ├── download_timm_pretrained.py              #   Download timm pretrained weights
│   ├── export_onnx.py                           #   ONNX model export (dynamic size + ORT verification)
│   ├── gen_standalone_yamls.py                  #   Generate standalone model YAML configs
│   ├── prepare_qata_mosmed.py                   #   QaTa-COV19 / MosMedData+ dataset validation
│   └── visualize.py                             #   Prediction visualization (input + pred + overlay)
├── docs/                         (61 docs)      # Detailed documentation
│   ├── tutorial/                 (31 files)     #   Step-by-step tutorial (01-09, 08a-08e sub-chapters, EN+CN, README, complete_guide)
│   ├── models/                                  #   Model docs: overview, networks, encoders, decoders, skip, bottleneck
│   ├── paradigms/                               #   Paradigm docs: infrastructure, semi, weak, DA, distillation, text-guided
│   ├── deployment/                              #   Deployment docs: ONNX, FLOPs, params, FPS
│   ├── data/                                    #   Data docs: 25 datasets, 5 types, 4 split modes
│   └── research_guide.md                        #   Research guide: 9 directions + 14 experiment scripts
├── train.py                                     # Supervised training (AMP + DDP + DataParallel + Logger + Warmup)
├── semi_train.py                                # Semi-supervised training (20 methods)
├── train_weakly_supervised.py                   # Weakly supervised training (20 methods)
├── train_domain_adaptation.py                   # Domain adaptation training (18 methods)
├── train_distillation.py                        # Knowledge distillation training (27 methods)
├── train_text_guided.py                         # Text-guided training (12 models)
├── test_text_guided.py                          # Text-guided inference (trainable + pipeline)
├── test.py                                      # Inference / testing
├── profile_model.py                             # FLOPs / params / FPS profiling
├── setup.py                                     # Package installation
└── requirements.txt                             # Python dependencies

🧩 Model Components

Detailed docs: docs/models/

Complete Networks — 130

CategoryCountExamples
CNN36UNet, UNet3+, UNet++, Attention-UNet, nnU-Net, MedNeXt, MEW-UNet, DCSAU-Net
Transformer36SegFormer, TransUNet, Swin-UNet, DAEFormer, MISSFormer, HiFormer, PolypPVT, CASCADE
Mamba / SSM24VM-UNet, U-Mamba, Swin-UMamba, LKM-UNet, LoG-VMamba, HC-Mamba
SAM family10MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, LiteMedSAM, SAMed, Medical SAM Adapter
KAN / MLP4RollingUNet, UNeXt, U-KAN, WA-UKAN
Linear Attention3TTT-UNet, U-VixLSTM, xLSTM-UNet
RWKV5U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV-UNet, RIR-Zigzag
Text-guided12CRIS, BiomedParse, LanGuideMedSeg, LViT, TGANet, TPRO, CausalCLIPSeg

Full list: docs/models/networks.md

Note on U-RWKV disambiguation: Two distinct networks share the "U-RWKV" name:

  • u_rwkvMICCAI 2025: Direction-Adaptive RWKV Module (DARM) + Stage-Adaptive Squeeze-and-Excitation (SASE), lightweight design with RWKV integrated within conv stages. Source: hbyecoding/U-RWKV
  • u_rwkv_tipIEEE TIP 2026: Standard U-Net + post-conv RWKV attention blocks with OmniShift multi-scale conv, originally for volumetric segmentation. Source: hbyecoding/U-RWKV

Encoders — 177

Highlight: 39 foundation model encoders covering 9 medical modalities

ModalityCountModels
General5DINOv2, DINOv3, DINO, CLIP-ViT, SAM-ViT
Pathology6Phikon, Phikon-v2, UNI, PLIP, MUSK, KEEP
Radiology4Rad-DINO, OmniRad, BioViL, CheXZero
Ophthalmology4RETFound-DINOv2, RETFound, FLAIR, OphMAE
Dermatology3DermCLIP, MoNet, PanDerm
General Medical3BiomedCLIP, MedCLIP, MedSigLIP
MLLM Vision8Qwen2.5-VL, Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, HealthGPT, HuLuMed, LingShu
Ultrasound3UltraFedFM, USF-MAE, SAMUS
Endoscopy3EndoViT, Endo-FM, Surgical-SAM

All foundation ViTs use DPT head (multi-block multi-scale features), not naive FPN-from-tokens.

Dynamic timm encoder: any model from timm.list_models() with timm_ prefix works directly.

encoder:
  name: timm_efficientnet_b7    # or any timm model name
  pretrained: true

Full list: docs/models/encoders.md

Decoders — 45

CategoryCountExamples
Basic (upsampling)4deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep
Dense (connections)2UNet++, UNet3+
Cascade10CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet
Attention6Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet
Transformer5DAEFormer, MISSFormer, MTUNet, SwinUNet, nnFormer
MLP2SegFormer MLP, MLP Decoder
Specific (network)13CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet
Mamba1VM-UNet
Pyramid2UPerNet, DeepLabV3 (ASPP)

Full list: docs/models/decoders.md

Skip Connections — docs/models/skip_connections.md

Bottlenecks — docs/models/bottlenecks.md


🎓 Training Paradigms

Detailed docs: docs/paradigms/

Infrastructure

FeatureYAML config
Mixed precision AMPtraining.amp: true or CLI --amp
Multi-GPU DDPtorchrun --nproc_per_node=N train.py
DataParalleltraining.parallel: dp
TensorBoardtraining.logger: tensorboard
WandBtraining.logger: wandb
Reproducibility Seedtraining.random_state: 42 + training.deterministic: true
Warmup schedulertraining.scheduler.name: warmup_cosine + warmup_epochs: 10
Config inheritance_base_: ../base.yaml
Albumentationstraining.augmentation: albumentations
YAML Aug Pipelinetraining.augmentation: pipeline + training.aug_pipeline: [...]

Full config guide: docs/paradigms/README.md

Augmentation Pipeline — 24 Methods

Freely combine 24 augmentation methods via YAML config, no code changes needed. All methods support intensity range parameters, randomly sampled per call.

training:
  augmentation: pipeline        # enable pipeline mode
  aug_pipeline:                 # define augmentations in order
    - name: horizontal_flip
      params: { p: 0.5 }
    - name: vertical_flip
      params: { p: 0.5 }
    - name: random_rotate90
      params: { p: 0.5 }
    - name: random_rotate
      params: { p: 0.3, degrees_range: [-30, 30] }
    - name: random_affine
      params: { p: 0.3, degrees_range: [-15, 15], translate_range: [0.0, 0.1], scale_range: [0.8, 1.2] }
    - name: elastic_deform
      params: { p: 0.3, alpha_range: [20, 80], sigma_range: [3, 7] }
    - name: copy_paste
      params: { p: 0.3, max_objects: 2, scale_range: [0.5, 1.5] }
    - name: mosaic
      params: { p: 0.3, offset_range: [0.0, 0.2] }
    - name: clahe
      params: { p: 0.3, clip_limit_range: [1.0, 5.0], tile_size_range: [4, 16] }
    - name: gamma_correction
      params: { p: 0.3, gamma_range: [0.7, 1.5] }
    - name: gaussian_blur
      params: { p: 0.2, kernel_range: [3, 7], sigma_range: [0.1, 2.0] }
    - name: gaussian_noise
      params: { p: 0.2, std_range: [0.01, 0.08] }

Supported Augmentation Methods (24):

CategoryMethods
Geometrichorizontal_flip, vertical_flip, random_rotate90, random_rotate, random_affine, random_perspective, random_scale, elastic_deform, grid_mask
Pixel-levelphotometric_distortion, color_jitter, brightness_contrast, gamma_correction, clahe, gaussian_blur, gaussian_noise, sharpness, posterize, random_solarize, channel_dropout
Maskingrandom_erasing, coarse_dropout, grid_mask
Sample-levelcopy_paste, mosaic

Note: All intensity parameters use _range suffix (e.g. degrees_range, alpha_range), randomly sampled per call.

Full parameter docs for each method: docs/data/README.md Full config example: resnet50_unet_advanced_aug.yaml

Semi-Supervised — 20 Methods

Mean Teacher · CPS · CCT · UniMatch · FixMatch · FlexMatch · FreeMatch · SoftMatch · UA-MT · URPC · Deep Co-Training · Pi-Model · Temporal Ensembling · Pseudo-Label · ICT · R-Drop · Cross-Teaching · CorrMatch · AllSpark · DiffRect

Details: docs/paradigms/semi_supervised.md

Domain Adaptation — 18 Methods

Source Only · AdvEnt · DANN · TENT · DPL · CBMT · FDA · CRST · PixMatch · MIC · DAFormer · HRDA · PiPa · DDB · SePiCo · DiGA · MICDrop · SemiVL

Details: docs/paradigms/domain_adaptation.md

Knowledge Distillation — 27 Methods

Vanilla KD · UNet-Distillation · FitNets · Attention Mimicry · AT · FSP · NST · RKD · VID · DKD · MGD · DIST · CIRKD · CWD · ReviewKD · SimKD · NORM · SDD · AICSD · LSKD · TTM · CTKD · MLKD + 4 medical-specific

Details: docs/paradigms/distillation.md

Weakly Supervised — 20 Methods

Box · CAM · MIL · Point · Scribble · TreeEnergy · SEAM · PuzzleCAM · AdvCAM · MCTformer · EPS · BoxInst · ReCAM · ToCo · LPCAM · MARS · DuPL · MoRe · PSDPM · SemPLeS

Details: docs/paradigms/weakly_supervised.md

Text-Guided — 12 Trainable Models + Inference Pipeline

Trainable models (12): CRIS · BiomedParse · LanGuideMedSeg · LViT · TGANet · TPRO · CausalCLIPSeg · CLIP-Universal · CXR-CLIP-Seg · TP-DRSeg · MedCLIP-SAM · SaLIP

Inference-only: MediSee (requires vendor model weights, see inference pipeline below)

Inference Pipeline (9 detector × 4 segmenter = 36 combinations):

  • Detector: GroundingDINO · Qwen2-VL · Qwen2.5-VL · Qwen3-VL · InternVL · LLaVA · MiniCPM-V · Phi3-V · CogVLM
  • Segmenter: SAM2 · MedSAM · SAM-Med2D · LiteMedSAM

Details: docs/paradigms/text_guided.md


⚡ Deployment & Efficiency

Detailed docs: docs/deployment/README.md

# ONNX Export
python scripts/export_onnx.py --config xxx.yaml --checkpoint best.pth --output model.onnx --verify

# FLOPs Calculation
python -c "
from fvcore.nn import FlopCountAnalysis
import torch
flops = FlopCountAnalysis(model, torch.randn(1,3,224,224))
print(f'FLOPs: {flops.total()/1e9:.2f}G')
"

# Params (trainable only)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable/1e6:.2f}M / Total: {total/1e6:.2f}M")

Note: Frozen foundation encoder params are NOT counted as trainable.


📊 Datasets

Detailed docs: docs/data/README.md Dataset example configs: configs/intro_to_datasets/

Supported Dataset Types

TypeDescription
synapseSynapse multi-organ CT (TransUNet format)
acdcACDC cardiac MRI (TransUNet format)
genericGeneric images/ + masks/ directories
qata_covid19QaTa-COV19 chest X-ray + per-image text (LViT format)
mosmed_plusMosMedData+ COVID CT + per-image text (LViT format)

Data Split Methods

# Method 1: Explicit paths
data:
  train_dir: ./data/train
  val_dir: ./data/val
  test_dir: ./data/test       # optional

# Method 2: Ratio-based split
data:
  root_dir: ./data/all
  train_ratio: 0.7
  val_ratio: 0.15

# Method 3: N-fold cross validation
data:
  root_dir: ./data/all
  n_splits: 5
  fold_idx: 0

Included Datasets (25)

CT: Synapse, COVID CT Seg, MosMedData+

MRI: ACDC

X-ray (CXR): Montgomery+Shenzhen, QaTa-COV19

Fundus Photography: DRIVE, STARE, CHASE_DB1, HRF, ARIA, RITE, REFUGE, Drishti-GS

Dermoscopy: ISIC 2016/2017/2018, PH2

Endoscopy: CVC-ClinicDB, CVC-ColonDB, Kvasir-SEG

Histopathology (WSI): GlaS, PanNuke, MoNuSeg

Ultrasound: BUSI


🔧 Config System

Two Model Config Modes

# Mode 1: Modular combination (encoder + decoder + skip + bottleneck)
model:
  num_classes: 9
  img_size: 224
  encoder:
    name: timm_resnet50
    pretrained: true
  decoder:
    name: unet
  skip_connection:
    name: concat
  bottleneck:
    name: aspp

# Mode 2: Complete architecture (architecture key)
model:
  num_classes: 9
  img_size: 224
  architecture: transunet
  arch_params: {}

Config Inheritance

# child.yaml — only write overrides
_base_: ../base_resnet50.yaml
model:
  num_classes: 9
training:
  epochs: 300

Complete Training Config Example

model:
  num_classes: 9
  img_size: 224
  transfer_learning_path: null   # full-model checkpoint for transfer learning
  encoder:
    name: timm_resnet50
    pretrained: true
    pretrained_path: null         # manual backbone checkpoint override
    in_channels: 3
  decoder:
    name: unet
    params: {}
  bottleneck:
    name: none

data:
  type: synapse
  img_size: 224
  train_dir: ./data/Synapse/train_npz
  val_dir: ./data/Synapse/test_vol_h5

training:
  random_state: 42
  deterministic: true
  amp: true
  parallel: auto
  logger: tensorboard
  augmentation: albumentations
  epochs: 200
  batch_size: 16
  num_workers: 4
  val_interval: 10
  loss:
    name: compound
    params:
      losses:
        - name: ce
          weight: 0.4
        - name: dice
          weight: 0.6
  optimizer:
    name: adamw
    lr: 0.0001
    weight_decay: 0.0001
  scheduler:
    name: warmup_cosine
    warmup_epochs: 10
    warmup_lr: 0.000001
    min_lr: 0.000001

🔌 Custom Extensions

Add New Encoder

# medseg/models/encoders/cnn/my_encoder.py
from medseg.registry import ENCODER_REGISTRY

@ENCODER_REGISTRY.register("my_encoder")
class MyEncoder(nn.Module):
    def __init__(self, pretrained=False, in_channels=3, img_size=224, **kwargs):
        super().__init__()
        self.out_channels = [64, 128, 256, 512]
    def forward(self, x):
        return [f1, f2, f3, f4]  # multi-scale features

Add New Decoder

@DECODER_REGISTRY.register("my_decoder")
class MyDecoder(nn.Module):
    has_internal_skip = False
    def __init__(self, encoder_channels, bottleneck_channels, skip_connection=None, **kwargs):
        super().__init__()
        self.out_channels = encoder_channels[0]
    def forward(self, bottleneck_feat, skip_features):
        return decoded

Add New Loss

@LOSS_REGISTRY.register("my_loss")
class MyLoss(nn.Module):
    def forward(self, pred, target):
        return loss_value

Add New Augmentation

# medseg/datasets/advanced_aug.py
from medseg.registry import AUGMENTATION_REGISTRY

@AUGMENTATION_REGISTRY.register("my_augmentation")
class MyAugmentation:
    def __init__(self, p=0.5, **kwargs):
        self.p = p

    def set_dataset(self, dataset):
        """Optional: implement if dataset access needed"""
        self.dataset = dataset

    def __call__(self, sample: dict) -> dict:
        import random
        if random.random() > self.p:
            return sample
        image, label = sample['image'], sample['label']
        # ... implement augmentation logic ...
        return {'image': image, 'label': label}

After registration and import in medseg/datasets/__init__.py, use via name: my_augmentation in YAML.

After registration and import in __init__.py, use via name: my_encoder in YAML.


📜 Citation & License

@misc{jiang2026aprilmedsegmodularmedicalimage,
      title={APRIL-MedSeg: A Modular Medical Image Segmentation Toolbox Embracing Modern Paradigms}, 
      author={Juntao Jiang and Jinsheng Bai and Linxuan Fan and Yali Bi and Jiangning Zhang and Yong Liu},
      year={2026},
      eprint={2606.30577},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.30577}, 
}

License

Apache 2.0. For legitimate academic research and engineering use only. Clinical deployment must comply with local regulations.

Acknowledgements

Thanks to PyTorch, timm, MONAI, SSL4MIS, SAM, GroundingDINO, DINOv2, CLIP, transformers, and all open-source projects that made this possible.


📬 Contact

For questions, collaborations, or bug reports, feel free to reach out:


🤝 Community

QQ Group
QQ Group QR Code
Scan to join QQ Group
Slack
Slack QR Code
Join our Slack workspace

Contributors

LightingMax

20 commits

LinLinLin-X

6 commits

juntaoJianggavin/APRIL-MedSeg

A Modern Modular 2D Medical Image Segmentation Toolbox

52

stars

45

commits

Python

primary language

Aug 14, 2026

updated

Browse cluster: Medical Image Segmentation with Foundation Models

README

APRIL-MedSeg Logo

License: Apache 2.0 arXiv

Juntao Jiang, Jinsheng Bai, Linxuan Fan, Yali Bi, Jiangning Zhang, Yong Liu

中文文档

130 networks · 177 encoders · 45 decoders · 81 losses · 25 skip connections · 17 bottlenecks · 6 training paradigms · 24 augmentations · 917 YAML configs · switch anything with one line of YAML


📰 Updates

  • 2026.07.30Bug fix: semi_train.py selected the best checkpoint with the EMA teacher but saved only the student weights, so best_model.pth could not reproduce the reported validation Dice for teacher-evaluated methods. The evaluated model is now stored in model_state_dict, and the student is kept in student_state_dict for resuming.
  • 2026.07.26 — Fixed RWKV-UNet WKV CUDA kernel loading and aligned GLSP with the official pretrained weights.
  • 2026.07.24Bug fix: Fixed multi-class mask loading across all datasets; pixel values are now correctly mapped to contiguous class indices.
  • 2026.07.20Bug fix: Fixed a critical data split issue in GenericDataset that caused incorrect train/val/test sample counts. All training/evaluation scripts updated. New models: Added support for SAM3 (Perception Encoder backbone), MedSAM2 (2D, Hiera backbone), and MedSAM3 (2D, PE backbone).
  • 2026.06.15 — The project has been officially renamed to APRIL-MedSeg, named after the APRIL Lab at Zhejiang University (led by Prof. Yong Liu).
  • 2026.06.11UltimateMedSeg is officially released!

📑 Table of Contents


📦 Installation

Requirements

  • Python >= 3.10
  • PyTorch >= 2.0
  • NVIDIA GPU with CUDA 12.4+

Basic Installation

git clone https://github.com/juntaoJianggavin/APRIL-MedSeg.git
cd APRIL-MedSeg

# Install dependencies
pip install -r requirements.txt

# Install in dev mode
pip install -e .

Optional Dependencies

# Foundation models
pip install transformers safetensors

# MLLM inference pipeline
pip install groundingdino-py
pip install git+https://github.com/facebookresearch/segment-anything.git

# ONNX export & verification
pip install onnx onnxruntime

# Lion optimizer
pip install lion-pytorch

# Mamba / SSM encoders (only required for Mamba-family networks; needs CUDA toolchain to build)
pip install causal-conv1d
pip install mamba-ssm

Pretrained Weights & Transfer Learning

Three levels of weight loading, from lightest to heaviest:

MethodYAML KeyScopeUse Case
Encoder pretrainedencoder.pretrained: trueBackbone onlyImageNet / domain-specific backbone weights
Manual pretrained pathencoder.pretrained_path: /path/to/weights.pthBackbone onlyOffline / custom backbone checkpoint
Transfer learningmodel.transfer_learning_path: /path/to/full_model.pthEntire networkLoad full model (encoder+decoder+bottleneck+head) from a previous training run
model:
  # Full-model transfer learning (loads AFTER encoder pretrained, takes priority)
  transfer_learning_path: null  # or /path/to/checkpoint.pth

  encoder:
    name: timm_resnet50
    pretrained: true              # auto-download ImageNet backbone weights
    pretrained_path: null         # or /path/to/backbone.pth (manual override)

Note: Models requiring specific pretrained weights (43 architectures in REQUIRES_PRETRAINED) will display a 10-second warning when pretrained: false. All auto-downloadable weights support manual path override via pretrained_path or transfer_learning_path.

Pretrained Weight Sources (4 Categories)

CategoryMechanismModels
A. WEIGHT_REGISTRY auto-downloadensure_weight() with GitHub/GCS/HF sourcesswinunet, h2former, hiformer, transunet, vm_unet, rwkv_unet (B/S/T), cswin_unet, da_transunet, mamba_unet, fcbformer, transnuseg
B. timm / torchvision runtimepretrained: true triggers built-in downloadsegformer_b0–b5, esfpnet, cascade, emcad, polyp_pvt, fatnet, transfuse, mist, hsnet, ssformer, ldnet, dconnnet, cfanet, lv_unet, nulite, polyper
C. SAM familypretrained: true auto-downloads ViT/SAM weightssam_b, sam_l, mobile_sam, sam2, sam_med2d, samed, sammed2d_wrapper, samus, auto_sam, lite_medsam, medical_sam_adapter
D. Foundation & MLLM encodersHuggingFace Hub / open_clip / transformers auto-download39 foundation encoders (dinov2, dino, clip_vit, sam_vit, dinov3, phikon, uni, plip, musk, phikon_v2, keep, raddino, omnirad, biovil, chexzero, retfound, retfound_dinov2, flair, ophmae, panderm, dermclip, monet_derm, endo_vit, endo_fm, surgical_sam, biomedclip, medclip, medsiglip, usfmae, ultrafedfm, samus) + 8 MLLM vision towers (qwen3_vl, qwen25_vl, llava_med, medgemma, healthgpt, huatuogpt, hulumed, lingshu)

Automatic Weight Download

# List all registered weights and cache status
python -m medseg.utils.weight_downloader list

# Download a specific weight (auto-retries all sources)
python -m medseg.utils.weight_downloader download medsam_vit_b

# Check which auto-downloadable weights are present
python -m medseg.utils.weight_downloader check

When all auto-download sources fail, the error message includes:

  • Manual download URL
  • Exact cache path to place the file at
  • Alternative: set pretrained_path in YAML to point to a local file

timm encoder weights are downloaded automatically via timm's built-in mechanism.


🚀 Quick Start

1. Standard Supervised Training

# ResNet50 + UNet decoder
python train.py --config configs/architectures/networks/general/aau_net.yaml \
    --output_dir output/aau_net

# With AMP mixed precision
python train.py --config configs/architectures/networks/general/transunet.yaml \
    --output_dir output/transunet --amp

# Multi-GPU DDP training
torchrun --nproc_per_node=4 train.py \
    --config configs/architectures/networks/general/swinunet.yaml \
    --output_dir output/swinunet --amp

2. Inference / Testing

# Single model evaluation
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth

# Save prediction results
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth --save_pred --output_dir test_output/

# Multi-checkpoint ensemble (logit averaging)
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint ckpt_a.pth ckpt_b.pth ckpt_c.pth \
    --ensemble-weights 0.5 0.3 0.2 \
    --ensemble-average logit

# Test-Time Augmentation (TTA)
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --tta \
    --tta-augs identity rot90 rot180 rot270 hflip vflip \
    --tta-merge mean

# TTA + Ensemble combined
python test.py --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint ckpt_a.pth ckpt_b.pth \
    --ensemble-average logit \
    --tta --tta-merge mean

3. Semi-Supervised Training

# Mean Teacher
python semi_train.py --config configs/training_paradigms/semi_supervision/mean_teacher.yaml \
    --output_dir output/semi_mt

# CPS (Cross Pseudo Supervision)
python semi_train.py --config configs/training_paradigms/semi_supervision/cps.yaml \
    --output_dir output/semi_cps

4. Domain Adaptation Training

# AdvEnt
python train_domain_adaptation.py \
    --config configs/training_paradigms/domain_adaptation/advent.yaml \
    --output_dir output/da_advent

# TENT (Test-Time Adaptation)
python train_domain_adaptation.py \
    --config configs/training_paradigms/domain_adaptation/tent.yaml \
    --output_dir output/da_tent

5. Knowledge Distillation Training

python train_distillation.py \
    --teacher_config configs/training_paradigms/distillation/teacher_large.yaml \
    --student_config configs/training_paradigms/distillation/student_small.yaml \
    --distillation_type logit \
    --temperature 4.0 \
    --alpha 0.5 \
    --output_dir output/kd_logit

6. Weakly Supervised Training

# Box-supervised
python train_weakly_supervised.py \
    --config configs/training_paradigms/weak_supervision/box_supervised.yaml \
    --supervision_type box \
    --output_dir output/weak_box

# CAM-based
python train_weakly_supervised.py \
    --config configs/training_paradigms/weak_supervision/cam.yaml \
    --supervision_type cam \
    --output_dir output/weak_cam

7. Text-Guided Training

# Train
python train_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_clip.yaml \
    --output_dir output/text_cris

# Test (auto-detects: trainable model vs inference pipeline)
python test_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_clip.yaml \
    --checkpoint output/text_cris/best_model.pth

# Test inference-only pipeline (no checkpoint needed)
python test_text_guided.py \
    --config configs/training_paradigms/text_guided/synapse_grounding_dino_sam2.yaml

8. Model Profiling

# FLOPs / Params / FPS
python profile_model.py --config configs/architectures/networks/general/transunet.yaml

9. ONNX Export

python scripts/export_onnx.py \
    --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --output model.onnx --verify

10. Prediction Visualization

python scripts/visualize.py \
    --config configs/architectures/networks/general/transunet.yaml \
    --checkpoint output/best_model.pth \
    --input ./data/test/images/ \
    --output vis_output/

11. Python API

from medseg.utils.config import load_config
from medseg.model_builder import build_model

cfg = load_config("configs/architectures/networks/general/transunet.yaml")
model = build_model(cfg)

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable params: {trainable / 1e6:.2f}M")

📚 Tutorial

A step-by-step tutorial series covering deep learning medical image segmentation from fundamentals to advanced topics:

ChapterTitleKey Topics
01Introduction to Medical Image SegmentationConcepts, clinical significance, metrics, method evolution
02U-Net in DetailArchitecture, skip connections, U-Net family variants
03Data and PreprocessingFormats, split strategies, augmentation pipeline
04Training and EvaluationLoss functions, optimizers, AMP/DDP, evaluation
05Encoder Deep DiveCNN / Transformer / Mamba / RWKV comparison, timm wrapper
06Decoders and Skip ConnectionsCASCADE / EMCAD / Attention Gate, skip taxonomy
07Foundation ModelsDPT head, 9 medical modalities, fine-tuning strategies
08Advanced Training ParadigmsOverview: semi-supervised, domain adaptation, distillation, weakly supervised
08aSemi-Supervised SegmentationMean Teacher, CPS, UniMatch, FixMatch, consistency regularization
08bDomain AdaptationAdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo
08cKnowledge DistillationVanillaKD, DKD, CWD, MGD, DIST, ReviewKD, SimKD
08dWeakly Supervised SegmentationCAM, SEAM, PuzzleCAM, Box, Point, Scribble supervision
08eText-Guided SegmentationCRIS, BiomedParse, LViT, CLIP-based, MLLM pipeline
09Deployment and InferenceONNX export, TTA, ensemble, MLLM pipeline

Full tutorial index


🏗️ Project Structure

segmentation_tool/
├── medseg/                                      # Core framework
│   ├── models/                                  # Model components
│   │   ├── encoders/                            #   177 encoders (92 native + 85 timm presets + 1000+ via timm_ prefix)
│   │   │   ├── cnn/              (13 modules)   #     CNN: basic, DCSAU, CFA, MedNeXt, MEW, R2U, AttUNet, LV, MALU, EGE, ConvNeXt, EfficientNetV2, HRNet
│   │   │   ├── transformer/      (18 modules)   #     Transformer: TransUNet, SwinUNet, MISSFormer, DAEFormer, HiFormer, PVTv2, MaxViT, ViT-Pyramid, ...
│   │   │   ├── mamba/            (10 modules)   #     Mamba/SSM: VMUNet, UMamba, LKM, LoG-VMamba, UltraLight-VM, VMKLA, ...
│   │   │   ├── rwkv/             (5 modules)    #     RWKV: RWKV-UNet, U-RWKV (MICCAI), U-RWKV (TIP), MD-RWKV, RIR-Zigzag
│   │   │   ├── linear_attn/      (5 modules)    #     Linear attention: RetNet, Linformer, Performer, TTT, xLSTM
│   │   │   ├── kan_mlp/          (4 modules)    #     KAN/MLP: UKAN, Rolling-UNet, UNeXt, WA-UKAN
│   │   │   ├── foundation/       (39 modules)   #     Foundation models (DPT head)
│   │   │   │   ├── general/      (5)            #       DINOv2, DINOv3, DINO, CLIP-ViT, SAM-ViT
│   │   │   │   ├── pathology/    (6)            #       Phikon, Phikon-v2, UNI, PLIP, MUSK, KEEP
│   │   │   │   ├── radiology/    (4)            #       Rad-DINO, OmniRad, BioViL, CheXZero
│   │   │   │   ├── ophthalmology/(4)            #       RETFound-DINOv2, FLAIR, OphMAE, RETFound
│   │   │   │   ├── dermatology/  (3)            #       PanDerm, DermCLIP, MonetDerm
│   │   │   │   ├── general_medical/(3)          #       BiomedCLIP, MedCLIP, MedSigLIP
│   │   │   │   ├── mllm_vision/  (8)            #       Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, ...
│   │   │   │   ├── endoscopy/    (3)            #       EndoViT, Endo-FM, Surgical-SAM
│   │   │   │   └── ultrasound/   (3)            #       UltraFedFM, USF-MAE, SAMUS
│   │   │   └── wrapper/          (1 module)     #     timm dynamic wrapper (85 pre-registered + 1000+ via timm_ prefix)
│   │   ├── decoders/                            #   45 decoders
│   │   │   ├── basic/            (4 registered) #     Basic upsampling: deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep
│   │   │   ├── dense/            (2 registered) #     Dense connections: UNet++, UNet3+
│   │   │   ├── cascade/          (10 registered)#     CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet
│   │   │   ├── attention/        (6 registered) #     Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet
│   │   │   ├── transformer/      (5 registered) #     DAEFormer, MISSFormer, MTUNet, nnFormer, SwinUNet
│   │   │   ├── mlp/              (2 registered) #     SegFormer MLP, MLP Decoder
│   │   │   ├── specific/         (13 registered)#     CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet
│   │   │   ├── pyramid/          (2 registered) #     UPerNet, DeepLabV3 (ASPP)
│   │   │   └── mamba/            (1 registered) #     VM-UNet
│   │   ├── bottlenecks/          (17 modules)   #   17 bottlenecks: none, basic, ASPP, DenseASPP, PPM, Transformer, SE, CBAM, ...
│   │   ├── skip_connections/                    #   25 skip connections
│   │   │   ├── basic/            (3 modules)    #     Basic: concat, dense, add
│   │   │   ├── attention/        (10 modules)   #     Attention: AG, CAB, SAB, SCSE, CBAM, Gating, GRU, GAB, SC-Att, TA-MoSC
│   │   │   ├── transformer/      (5 modules)    #     Transformer: CrossAttn, TransFusion, AggAttn, MISSFormer, UCTrans
│   │   │   ├── mamba/            (1 module)     #     Mamba: SK-VM++
│   │   │   └── fusion/           (6 modules)    #     CNN fusion: BiFusion, Deformable, MultiScale, FeatureRefine, CCM, SDI
│   │   ├── networks/                            #   130 complete architectures (variants merged)
│   │   │   ├── cnn/              (36 registered)#     CNN: UNet, UNet3+, UNet++, AttUNet, nnUNet, MedNeXt, MEW-UNet, ...
│   │   │   ├── transformer/      (36 registered)#     Transformer: SegFormer, TransUNet, SwinUNet, DAEFormer, PolypPVT, CASCADE, ...
│   │   │   ├── mamba/            (24 registered)#     Mamba: VMUNet, U-Mamba, SwinUMamba, SkinMamba, DermoMamba, SerpMamba, ...
│   │   │   ├── sam/              (10 registered)#     SAM family: MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, ...
│   │   │   ├── rwkv/             (5 registered) #     RWKV: U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV, RIR-Zigzag
│   │   │   ├── kan_mlp/          (4 registered) #     KAN/MLP: RollingUNet, UNeXt, UKAN, WA-UKAN
│   │   │   └── linear_attn/      (3 registered) #     Linear attention: TTT-UNet, U-VixLSTM, xLSTM-UNet
│   │   └── text_unet/            (13 modules)   #   Text-guided (12 models): CRIS, BiomedParse, LanGuideMedSeg, LViT, TGANet, TPRO, ...
│   ├── training/                                # Training paradigms
│   │   ├── semi/                 (22 modules)   #   20 semi-supervised methods + 2 utils (base, utils)
│   │   │                                        #     MeanTeacher, CPS, UniMatch, FixMatch, UA-MT, CorrMatch, AllSpark, ...
│   │   ├── domain_adaptation/    (18 modules)   #   18 domain adaptation: AdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo, ...
│   │   ├── distillation/         (28 modules)   #   27 distillation: VanillaKD, DKD, MGD, DIST, CWD, ReviewKD, SimKD, NORM, ...
│   │   └── weakly_supervised/    (21 modules)   #   20 weakly supervised methods (CAM, SEAM, PuzzleCAM, TreeEnergy, ...)
│   ├── inference/                               # Inference
│   │   ├── ensemble.py                          #   Ensemble inference (multi-model voting)
│   │   ├── tta.py                               #   Test-time augmentation
│   │   └── mllm/                 (16 modules)   #   MLLM pipeline: 9 detector × 4 segmenter = 36 combinations
│   │       │                                    #     Detector: GroundingDINO, Qwen2/2.5/3-VL, InternVL, LLaVA, MiniCPM-V, Phi3-V, CogVLM
│   │       │                                    #     Segmenter: SAM2, MedSAM, SAM-Med2D, LiteMedSAM
│   │       └── medisee/          (3 modules)    #     MediSee: LLM reasoning segmenter
│   ├── losses/                   (15 modules)   # 81 losses
│   │                                            #   Supervised: CE, Dice, Focal, Tversky, Lovász, Boundary, Hausdorff, ...
│   │                                            #   Distillation: VanillaKD, DKD, CWD, MGD, DIST, AT, RKD, ...
│   │                                            #   Domain adaptation: AdvEnt, DANN, FDA, MIC, TENT, ...
│   │                                            #   Weakly supervised: Box, CAM, Point, Scribble, TreeEnergy, SEAM, ...
│   ├── datasets/                 (10 modules)   # Data loading: Synapse, ACDC, Generic, QaTa-COV19, MosMedData+, 24 augmentations
│   │   ├── advanced_aug.py                      #   24 advanced augmentations (YAML configurable)
│   │   └── transforms.py                        #   Basic transforms (Resize, ToTensor, Normalize)
│   ├── utils/                    (11 modules)   # Utilities
│   │   ├── amp_ddp.py                           #   AMP mixed precision + DDP distributed + DataParallel
│   │   ├── logger.py                            #   TensorBoard / WandB unified logging
│   │   ├── config.py                            #   Config inheritance (_base_ field support)
│   │   ├── warmup.py                            #   Warmup scheduler + Lion/AdamW/SGD optimizers
│   │   ├── augmentation.py                      #   Augmentation builder (basic/albumentations/pipeline)
│   │   ├── reproducibility.py                   #   Reproducibility (global seed + cuDNN deterministic)
│   │   ├── weight_downloader.py                 #   Automatic weight download + manual URL hints
│   │   ├── metrics.py                           #   Evaluation metrics: Dice, IoU, HD95, NSD
│   │   ├── hf_hub.py                            #   HuggingFace Hub model/dataset download
│   │   ├── timm_compat.py                       #   timm version compatibility utilities
│   │   └── timm_pretrained.py                   #   timm pretrained weight management
│   ├── text_guided.py                           # Text-guided segmentation (CRIS, BiomedParse, LanGuideMedSeg, ...)
│   ├── model_builder.py                         # YAML → model auto-assembler
│   └── registry.py                              # 6 registries: ENCODER / DECODER / SKIP / BOTTLENECK / LOSS / AUGMENTATION
├── data/                                        # Dataset root (user datasets go here)
│   ├── YourDataset/                             #   Your custom dataset
│   ├── source/                                  #   Domain adaptation source
│   ├── target/                                  #   Domain adaptation target
│   ├── target_val/                              #   Domain adaptation validation
│   └── test_dummy/                              #   Dummy test data
├── figs/                                        # Figures & logos
│   └── logo.png                                 #   Project logo
├── configs/                      (917 yamls)    # YAML configs
│   ├── architectures/            (783 yamls)    #   Network architecture configs
│   │   ├── networks/             (302 yamls)    #     Complete networks (130 arch across general/acdc/synapse)
│   │   ├── combinations/         (169 yamls)    #     Encoder+decoder free combinations
│   │   ├── decoder_study/        (133 yamls)    #     Decoder ablation (3 enc × 44 dec + 1)
│   │   ├── skip_study/           (75 yamls)     #     Skip ablation (3 enc × 25 skip)
│   │   ├── bottleneck_study/     (51 yamls)     #     Bottleneck ablation (3 enc × 17 bn)
│   │   └── foundation/           (53 yamls)     #     Foundation models (9 modalities × 39 encoders)
│   ├── training_paradigms/       (104 yamls)    #   Training paradigm configs
│   │   ├── semi_supervision/     (20 yamls)     #     Semi-supervised (20 methods)
│   │   ├── domain_adaptation/    (18 yamls)     #     Domain adaptation (18 methods)
│   │   ├── distillation/         (29 yamls)     #     Distillation (27 methods)
│   │   ├── text_guided/          (17 yamls)     #     Text-guided (12 models + pipeline)
│   │   └── weak_supervision/     (20 yamls)     #     Weakly supervised (20 methods)
│   └── intro_to_datasets/        (27 yamls)     #   27 dataset introductions + example configs
├── scripts/                                     # Utility + experiment scripts
│   ├── experiments/              (14 scripts)   #   Experiment bash scripts
│   │   ├── run_sota_benchmark.sh                #     SOTA architecture comparison (11 models × 7 datasets)
│   │   ├── run_decoder_study.sh                 #     Decoder ablation (3 enc × 15 classic dec)
│   │   ├── run_bottleneck_study.sh              #     Bottleneck ablation (3 enc × 9 bn)
│   │   ├── run_skip_study.sh                    #     Skip ablation (3 enc × 12 skip)
│   │   ├── run_polyp_benchmark.sh               #     Polyp-specific models (16 models × 2 datasets)
│   │   ├── run_skin_benchmark.sh                #     Skin-specific models (16 models × 2 datasets + PH2 external)
│   │   ├── run_retinal_benchmark.sh             #     Retinal-specific models (7 models × 3 datasets)
│   │   ├── run_ultrasound_benchmark.sh          #     Ultrasound-specific models (8 models × BUSI)
│   │   ├── run_pathology_benchmark.sh           #     Pathology-specific models (5 models × GlaS)
│   │   ├── run_lightweight_skin.sh              #     Lightweight skin segmentation (8 models)
│   │   ├── run_semi_study.sh                    #     Semi-supervised paradigm comparison (6 methods)
│   │   ├── run_da_study.sh                      #     Domain adaptation paradigm comparison (8 methods)
│   │   ├── run_kd_study.sh                      #     Knowledge distillation comparison (7 methods)
│   │   └── run_weak_study.sh                    #     Weakly supervised paradigm comparison (6 methods)
│   ├── check_config_paths.py                    #   Check config path references across docs/scripts
│   ├── download_hf_dataset.py                   #   Download HuggingFace datasets
│   ├── download_timm_pretrained.py              #   Download timm pretrained weights
│   ├── export_onnx.py                           #   ONNX model export (dynamic size + ORT verification)
│   ├── gen_standalone_yamls.py                  #   Generate standalone model YAML configs
│   ├── prepare_qata_mosmed.py                   #   QaTa-COV19 / MosMedData+ dataset validation
│   └── visualize.py                             #   Prediction visualization (input + pred + overlay)
├── docs/                         (61 docs)      # Detailed documentation
│   ├── tutorial/                 (31 files)     #   Step-by-step tutorial (01-09, 08a-08e sub-chapters, EN+CN, README, complete_guide)
│   ├── models/                                  #   Model docs: overview, networks, encoders, decoders, skip, bottleneck
│   ├── paradigms/                               #   Paradigm docs: infrastructure, semi, weak, DA, distillation, text-guided
│   ├── deployment/                              #   Deployment docs: ONNX, FLOPs, params, FPS
│   ├── data/                                    #   Data docs: 25 datasets, 5 types, 4 split modes
│   └── research_guide.md                        #   Research guide: 9 directions + 14 experiment scripts
├── train.py                                     # Supervised training (AMP + DDP + DataParallel + Logger + Warmup)
├── semi_train.py                                # Semi-supervised training (20 methods)
├── train_weakly_supervised.py                   # Weakly supervised training (20 methods)
├── train_domain_adaptation.py                   # Domain adaptation training (18 methods)
├── train_distillation.py                        # Knowledge distillation training (27 methods)
├── train_text_guided.py                         # Text-guided training (12 models)
├── test_text_guided.py                          # Text-guided inference (trainable + pipeline)
├── test.py                                      # Inference / testing
├── profile_model.py                             # FLOPs / params / FPS profiling
├── setup.py                                     # Package installation
└── requirements.txt                             # Python dependencies

🧩 Model Components

Detailed docs: docs/models/

Complete Networks — 130

CategoryCountExamples
CNN36UNet, UNet3+, UNet++, Attention-UNet, nnU-Net, MedNeXt, MEW-UNet, DCSAU-Net
Transformer36SegFormer, TransUNet, Swin-UNet, DAEFormer, MISSFormer, HiFormer, PolypPVT, CASCADE
Mamba / SSM24VM-UNet, U-Mamba, Swin-UMamba, LKM-UNet, LoG-VMamba, HC-Mamba
SAM family10MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, LiteMedSAM, SAMed, Medical SAM Adapter
KAN / MLP4RollingUNet, UNeXt, U-KAN, WA-UKAN
Linear Attention3TTT-UNet, U-VixLSTM, xLSTM-UNet
RWKV5U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV-UNet, RIR-Zigzag
Text-guided12CRIS, BiomedParse, LanGuideMedSeg, LViT, TGANet, TPRO, CausalCLIPSeg

Full list: docs/models/networks.md

Note on U-RWKV disambiguation: Two distinct networks share the "U-RWKV" name:

  • u_rwkvMICCAI 2025: Direction-Adaptive RWKV Module (DARM) + Stage-Adaptive Squeeze-and-Excitation (SASE), lightweight design with RWKV integrated within conv stages. Source: hbyecoding/U-RWKV
  • u_rwkv_tipIEEE TIP 2026: Standard U-Net + post-conv RWKV attention blocks with OmniShift multi-scale conv, originally for volumetric segmentation. Source: hbyecoding/U-RWKV

Encoders — 177

Highlight: 39 foundation model encoders covering 9 medical modalities

ModalityCountModels
General5DINOv2, DINOv3, DINO, CLIP-ViT, SAM-ViT
Pathology6Phikon, Phikon-v2, UNI, PLIP, MUSK, KEEP
Radiology4Rad-DINO, OmniRad, BioViL, CheXZero
Ophthalmology4RETFound-DINOv2, RETFound, FLAIR, OphMAE
Dermatology3DermCLIP, MoNet, PanDerm
General Medical3BiomedCLIP, MedCLIP, MedSigLIP
MLLM Vision8Qwen2.5-VL, Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, HealthGPT, HuLuMed, LingShu
Ultrasound3UltraFedFM, USF-MAE, SAMUS
Endoscopy3EndoViT, Endo-FM, Surgical-SAM

All foundation ViTs use DPT head (multi-block multi-scale features), not naive FPN-from-tokens.

Dynamic timm encoder: any model from timm.list_models() with timm_ prefix works directly.

encoder:
  name: timm_efficientnet_b7    # or any timm model name
  pretrained: true

Full list: docs/models/encoders.md

Decoders — 45

CategoryCountExamples
Basic (upsampling)4deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep
Dense (connections)2UNet++, UNet3+
Cascade10CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet
Attention6Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet
Transformer5DAEFormer, MISSFormer, MTUNet, SwinUNet, nnFormer
MLP2SegFormer MLP, MLP Decoder
Specific (network)13CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet
Mamba1VM-UNet
Pyramid2UPerNet, DeepLabV3 (ASPP)

Full list: docs/models/decoders.md

Skip Connections — docs/models/skip_connections.md

Bottlenecks — docs/models/bottlenecks.md


🎓 Training Paradigms

Detailed docs: docs/paradigms/

Infrastructure

FeatureYAML config
Mixed precision AMPtraining.amp: true or CLI --amp
Multi-GPU DDPtorchrun --nproc_per_node=N train.py
DataParalleltraining.parallel: dp
TensorBoardtraining.logger: tensorboard
WandBtraining.logger: wandb
Reproducibility Seedtraining.random_state: 42 + training.deterministic: true
Warmup schedulertraining.scheduler.name: warmup_cosine + warmup_epochs: 10
Config inheritance_base_: ../base.yaml
Albumentationstraining.augmentation: albumentations
YAML Aug Pipelinetraining.augmentation: pipeline + training.aug_pipeline: [...]

Full config guide: docs/paradigms/README.md

Augmentation Pipeline — 24 Methods

Freely combine 24 augmentation methods via YAML config, no code changes needed. All methods support intensity range parameters, randomly sampled per call.

training:
  augmentation: pipeline        # enable pipeline mode
  aug_pipeline:                 # define augmentations in order
    - name: horizontal_flip
      params: { p: 0.5 }
    - name: vertical_flip
      params: { p: 0.5 }
    - name: random_rotate90
      params: { p: 0.5 }
    - name: random_rotate
      params: { p: 0.3, degrees_range: [-30, 30] }
    - name: random_affine
      params: { p: 0.3, degrees_range: [-15, 15], translate_range: [0.0, 0.1], scale_range: [0.8, 1.2] }
    - name: elastic_deform
      params: { p: 0.3, alpha_range: [20, 80], sigma_range: [3, 7] }
    - name: copy_paste
      params: { p: 0.3, max_objects: 2, scale_range: [0.5, 1.5] }
    - name: mosaic
      params: { p: 0.3, offset_range: [0.0, 0.2] }
    - name: clahe
      params: { p: 0.3, clip_limit_range: [1.0, 5.0], tile_size_range: [4, 16] }
    - name: gamma_correction
      params: { p: 0.3, gamma_range: [0.7, 1.5] }
    - name: gaussian_blur
      params: { p: 0.2, kernel_range: [3, 7], sigma_range: [0.1, 2.0] }
    - name: gaussian_noise
      params: { p: 0.2, std_range: [0.01, 0.08] }

Supported Augmentation Methods (24):

CategoryMethods
Geometrichorizontal_flip, vertical_flip, random_rotate90, random_rotate, random_affine, random_perspective, random_scale, elastic_deform, grid_mask
Pixel-levelphotometric_distortion, color_jitter, brightness_contrast, gamma_correction, clahe, gaussian_blur, gaussian_noise, sharpness, posterize, random_solarize, channel_dropout
Maskingrandom_erasing, coarse_dropout, grid_mask
Sample-levelcopy_paste, mosaic

Note: All intensity parameters use _range suffix (e.g. degrees_range, alpha_range), randomly sampled per call.

Full parameter docs for each method: docs/data/README.md Full config example: resnet50_unet_advanced_aug.yaml

Semi-Supervised — 20 Methods

Mean Teacher · CPS · CCT · UniMatch · FixMatch · FlexMatch · FreeMatch · SoftMatch · UA-MT · URPC · Deep Co-Training · Pi-Model · Temporal Ensembling · Pseudo-Label · ICT · R-Drop · Cross-Teaching · CorrMatch · AllSpark · DiffRect

Details: docs/paradigms/semi_supervised.md

Domain Adaptation — 18 Methods

Source Only · AdvEnt · DANN · TENT · DPL · CBMT · FDA · CRST · PixMatch · MIC · DAFormer · HRDA · PiPa · DDB · SePiCo · DiGA · MICDrop · SemiVL

Details: docs/paradigms/domain_adaptation.md

Knowledge Distillation — 27 Methods

Vanilla KD · UNet-Distillation · FitNets · Attention Mimicry · AT · FSP · NST · RKD · VID · DKD · MGD · DIST · CIRKD · CWD · ReviewKD · SimKD · NORM · SDD · AICSD · LSKD · TTM · CTKD · MLKD + 4 medical-specific

Details: docs/paradigms/distillation.md

Weakly Supervised — 20 Methods

Box · CAM · MIL · Point · Scribble · TreeEnergy · SEAM · PuzzleCAM · AdvCAM · MCTformer · EPS · BoxInst · ReCAM · ToCo · LPCAM · MARS · DuPL · MoRe · PSDPM · SemPLeS

Details: docs/paradigms/weakly_supervised.md

Text-Guided — 12 Trainable Models + Inference Pipeline

Trainable models (12): CRIS · BiomedParse · LanGuideMedSeg · LViT · TGANet · TPRO · CausalCLIPSeg · CLIP-Universal · CXR-CLIP-Seg · TP-DRSeg · MedCLIP-SAM · SaLIP

Inference-only: MediSee (requires vendor model weights, see inference pipeline below)

Inference Pipeline (9 detector × 4 segmenter = 36 combinations):

  • Detector: GroundingDINO · Qwen2-VL · Qwen2.5-VL · Qwen3-VL · InternVL · LLaVA · MiniCPM-V · Phi3-V · CogVLM
  • Segmenter: SAM2 · MedSAM · SAM-Med2D · LiteMedSAM

Details: docs/paradigms/text_guided.md


⚡ Deployment & Efficiency

Detailed docs: docs/deployment/README.md

# ONNX Export
python scripts/export_onnx.py --config xxx.yaml --checkpoint best.pth --output model.onnx --verify

# FLOPs Calculation
python -c "
from fvcore.nn import FlopCountAnalysis
import torch
flops = FlopCountAnalysis(model, torch.randn(1,3,224,224))
print(f'FLOPs: {flops.total()/1e9:.2f}G')
"

# Params (trainable only)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable/1e6:.2f}M / Total: {total/1e6:.2f}M")

Note: Frozen foundation encoder params are NOT counted as trainable.


📊 Datasets

Detailed docs: docs/data/README.md Dataset example configs: configs/intro_to_datasets/

Supported Dataset Types

TypeDescription
synapseSynapse multi-organ CT (TransUNet format)
acdcACDC cardiac MRI (TransUNet format)
genericGeneric images/ + masks/ directories
qata_covid19QaTa-COV19 chest X-ray + per-image text (LViT format)
mosmed_plusMosMedData+ COVID CT + per-image text (LViT format)

Data Split Methods

# Method 1: Explicit paths
data:
  train_dir: ./data/train
  val_dir: ./data/val
  test_dir: ./data/test       # optional

# Method 2: Ratio-based split
data:
  root_dir: ./data/all
  train_ratio: 0.7
  val_ratio: 0.15

# Method 3: N-fold cross validation
data:
  root_dir: ./data/all
  n_splits: 5
  fold_idx: 0

Included Datasets (25)

CT: Synapse, COVID CT Seg, MosMedData+

MRI: ACDC

X-ray (CXR): Montgomery+Shenzhen, QaTa-COV19

Fundus Photography: DRIVE, STARE, CHASE_DB1, HRF, ARIA, RITE, REFUGE, Drishti-GS

Dermoscopy: ISIC 2016/2017/2018, PH2

Endoscopy: CVC-ClinicDB, CVC-ColonDB, Kvasir-SEG

Histopathology (WSI): GlaS, PanNuke, MoNuSeg

Ultrasound: BUSI


🔧 Config System

Two Model Config Modes

# Mode 1: Modular combination (encoder + decoder + skip + bottleneck)
model:
  num_classes: 9
  img_size: 224
  encoder:
    name: timm_resnet50
    pretrained: true
  decoder:
    name: unet
  skip_connection:
    name: concat
  bottleneck:
    name: aspp

# Mode 2: Complete architecture (architecture key)
model:
  num_classes: 9
  img_size: 224
  architecture: transunet
  arch_params: {}

Config Inheritance

# child.yaml — only write overrides
_base_: ../base_resnet50.yaml
model:
  num_classes: 9
training:
  epochs: 300

Complete Training Config Example

model:
  num_classes: 9
  img_size: 224
  transfer_learning_path: null   # full-model checkpoint for transfer learning
  encoder:
    name: timm_resnet50
    pretrained: true
    pretrained_path: null         # manual backbone checkpoint override
    in_channels: 3
  decoder:
    name: unet
    params: {}
  bottleneck:
    name: none

data:
  type: synapse
  img_size: 224
  train_dir: ./data/Synapse/train_npz
  val_dir: ./data/Synapse/test_vol_h5

training:
  random_state: 42
  deterministic: true
  amp: true
  parallel: auto
  logger: tensorboard
  augmentation: albumentations
  epochs: 200
  batch_size: 16
  num_workers: 4
  val_interval: 10
  loss:
    name: compound
    params:
      losses:
        - name: ce
          weight: 0.4
        - name: dice
          weight: 0.6
  optimizer:
    name: adamw
    lr: 0.0001
    weight_decay: 0.0001
  scheduler:
    name: warmup_cosine
    warmup_epochs: 10
    warmup_lr: 0.000001
    min_lr: 0.000001

🔌 Custom Extensions

Add New Encoder

# medseg/models/encoders/cnn/my_encoder.py
from medseg.registry import ENCODER_REGISTRY

@ENCODER_REGISTRY.register("my_encoder")
class MyEncoder(nn.Module):
    def __init__(self, pretrained=False, in_channels=3, img_size=224, **kwargs):
        super().__init__()
        self.out_channels = [64, 128, 256, 512]
    def forward(self, x):
        return [f1, f2, f3, f4]  # multi-scale features

Add New Decoder

@DECODER_REGISTRY.register("my_decoder")
class MyDecoder(nn.Module):
    has_internal_skip = False
    def __init__(self, encoder_channels, bottleneck_channels, skip_connection=None, **kwargs):
        super().__init__()
        self.out_channels = encoder_channels[0]
    def forward(self, bottleneck_feat, skip_features):
        return decoded

Add New Loss

@LOSS_REGISTRY.register("my_loss")
class MyLoss(nn.Module):
    def forward(self, pred, target):
        return loss_value

Add New Augmentation

# medseg/datasets/advanced_aug.py
from medseg.registry import AUGMENTATION_REGISTRY

@AUGMENTATION_REGISTRY.register("my_augmentation")
class MyAugmentation:
    def __init__(self, p=0.5, **kwargs):
        self.p = p

    def set_dataset(self, dataset):
        """Optional: implement if dataset access needed"""
        self.dataset = dataset

    def __call__(self, sample: dict) -> dict:
        import random
        if random.random() > self.p:
            return sample
        image, label = sample['image'], sample['label']
        # ... implement augmentation logic ...
        return {'image': image, 'label': label}

After registration and import in medseg/datasets/__init__.py, use via name: my_augmentation in YAML.

After registration and import in __init__.py, use via name: my_encoder in YAML.


📜 Citation & License

@misc{jiang2026aprilmedsegmodularmedicalimage,
      title={APRIL-MedSeg: A Modular Medical Image Segmentation Toolbox Embracing Modern Paradigms}, 
      author={Juntao Jiang and Jinsheng Bai and Linxuan Fan and Yali Bi and Jiangning Zhang and Yong Liu},
      year={2026},
      eprint={2606.30577},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.30577}, 
}

License

Apache 2.0. For legitimate academic research and engineering use only. Clinical deployment must comply with local regulations.

Acknowledgements

Thanks to PyTorch, timm, MONAI, SSL4MIS, SAM, GroundingDINO, DINOv2, CLIP, transformers, and all open-source projects that made this possible.


📬 Contact

For questions, collaborations, or bug reports, feel free to reach out:


🤝 Community

QQ Group
QQ Group QR Code
Scan to join QQ Group
Slack
Slack QR Code
Join our Slack workspace

Contributors

LightingMax

20 commits

LinLinLin-X

6 commits

Languages

Python

99.3%