A Modern Modular 2D Medical Image Segmentation Toolbox
52
stars
45
commits
Python
primary language
Aug 14, 2026
updated
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
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.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).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 .
# 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
Three levels of weight loading, from lightest to heaviest:
| Method | YAML Key | Scope | Use Case |
|---|---|---|---|
| Encoder pretrained | encoder.pretrained: true | Backbone only | ImageNet / domain-specific backbone weights |
| Manual pretrained path | encoder.pretrained_path: /path/to/weights.pth | Backbone only | Offline / custom backbone checkpoint |
| Transfer learning | model.transfer_learning_path: /path/to/full_model.pth | Entire network | Load 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 whenpretrained: false. All auto-downloadable weights support manual path override viapretrained_pathortransfer_learning_path.
| Category | Mechanism | Models |
|---|---|---|
| A. WEIGHT_REGISTRY auto-download | ensure_weight() with GitHub/GCS/HF sources | swinunet, h2former, hiformer, transunet, vm_unet, rwkv_unet (B/S/T), cswin_unet, da_transunet, mamba_unet, fcbformer, transnuseg |
| B. timm / torchvision runtime | pretrained: true triggers built-in download | segformer_b0–b5, esfpnet, cascade, emcad, polyp_pvt, fatnet, transfuse, mist, hsnet, ssformer, ldnet, dconnnet, cfanet, lv_unet, nulite, polyper |
| C. SAM family | pretrained: true auto-downloads ViT/SAM weights | sam_b, sam_l, mobile_sam, sam2, sam_med2d, samed, sammed2d_wrapper, samus, auto_sam, lite_medsam, medical_sam_adapter |
| D. Foundation & MLLM encoders | HuggingFace Hub / open_clip / transformers auto-download | 39 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) |
# 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:
pretrained_path in YAML to point to a local filetimm encoder weights are downloaded automatically via timm's built-in mechanism.
# 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
# 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
# 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
# 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
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
# 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
# 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
# FLOPs / Params / FPS
python profile_model.py --config configs/architectures/networks/general/transunet.yaml
python scripts/export_onnx.py \
--config configs/architectures/networks/general/transunet.yaml \
--checkpoint output/best_model.pth \
--output model.onnx --verify
python scripts/visualize.py \
--config configs/architectures/networks/general/transunet.yaml \
--checkpoint output/best_model.pth \
--input ./data/test/images/ \
--output vis_output/
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")
A step-by-step tutorial series covering deep learning medical image segmentation from fundamentals to advanced topics:
| Chapter | Title | Key Topics |
|---|---|---|
| 01 | Introduction to Medical Image Segmentation | Concepts, clinical significance, metrics, method evolution |
| 02 | U-Net in Detail | Architecture, skip connections, U-Net family variants |
| 03 | Data and Preprocessing | Formats, split strategies, augmentation pipeline |
| 04 | Training and Evaluation | Loss functions, optimizers, AMP/DDP, evaluation |
| 05 | Encoder Deep Dive | CNN / Transformer / Mamba / RWKV comparison, timm wrapper |
| 06 | Decoders and Skip Connections | CASCADE / EMCAD / Attention Gate, skip taxonomy |
| 07 | Foundation Models | DPT head, 9 medical modalities, fine-tuning strategies |
| 08 | Advanced Training Paradigms | Overview: semi-supervised, domain adaptation, distillation, weakly supervised |
| 08a | Semi-Supervised Segmentation | Mean Teacher, CPS, UniMatch, FixMatch, consistency regularization |
| 08b | Domain Adaptation | AdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo |
| 08c | Knowledge Distillation | VanillaKD, DKD, CWD, MGD, DIST, ReviewKD, SimKD |
| 08d | Weakly Supervised Segmentation | CAM, SEAM, PuzzleCAM, Box, Point, Scribble supervision |
| 08e | Text-Guided Segmentation | CRIS, BiomedParse, LViT, CLIP-based, MLLM pipeline |
| 09 | Deployment and Inference | ONNX export, TTA, ensemble, MLLM pipeline |
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
Detailed docs: docs/models/
| Category | Count | Examples |
|---|---|---|
| CNN | 36 | UNet, UNet3+, UNet++, Attention-UNet, nnU-Net, MedNeXt, MEW-UNet, DCSAU-Net |
| Transformer | 36 | SegFormer, TransUNet, Swin-UNet, DAEFormer, MISSFormer, HiFormer, PolypPVT, CASCADE |
| Mamba / SSM | 24 | VM-UNet, U-Mamba, Swin-UMamba, LKM-UNet, LoG-VMamba, HC-Mamba |
| SAM family | 10 | MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, LiteMedSAM, SAMed, Medical SAM Adapter |
| KAN / MLP | 4 | RollingUNet, UNeXt, U-KAN, WA-UKAN |
| Linear Attention | 3 | TTT-UNet, U-VixLSTM, xLSTM-UNet |
| RWKV | 5 | U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV-UNet, RIR-Zigzag |
| Text-guided | 12 | CRIS, 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_rwkv— MICCAI 2025: Direction-Adaptive RWKV Module (DARM) + Stage-Adaptive Squeeze-and-Excitation (SASE), lightweight design with RWKV integrated within conv stages. Source: hbyecoding/U-RWKVu_rwkv_tip— IEEE TIP 2026: Standard U-Net + post-conv RWKV attention blocks with OmniShift multi-scale conv, originally for volumetric segmentation. Source: hbyecoding/U-RWKV
Highlight: 39 foundation model encoders covering 9 medical modalities
| Modality | Count | Models |
|---|---|---|
| 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, RETFound, FLAIR, OphMAE |
| Dermatology | 3 | DermCLIP, MoNet, PanDerm |
| General Medical | 3 | BiomedCLIP, MedCLIP, MedSigLIP |
| MLLM Vision | 8 | Qwen2.5-VL, Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, HealthGPT, HuLuMed, LingShu |
| Ultrasound | 3 | UltraFedFM, USF-MAE, SAMUS |
| Endoscopy | 3 | EndoViT, 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
| Category | Count | Examples |
|---|---|---|
| Basic (upsampling) | 4 | deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep |
| Dense (connections) | 2 | UNet++, UNet3+ |
| Cascade | 10 | CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet |
| Attention | 6 | Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet |
| Transformer | 5 | DAEFormer, MISSFormer, MTUNet, SwinUNet, nnFormer |
| MLP | 2 | SegFormer MLP, MLP Decoder |
| Specific (network) | 13 | CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet |
| Mamba | 1 | VM-UNet |
| Pyramid | 2 | UPerNet, DeepLabV3 (ASPP) |
Full list: docs/models/decoders.md
Detailed docs: docs/paradigms/
| Feature | YAML config |
|---|---|
| Mixed precision AMP | training.amp: true or CLI --amp |
| Multi-GPU DDP | torchrun --nproc_per_node=N train.py |
| DataParallel | training.parallel: dp |
| TensorBoard | training.logger: tensorboard |
| WandB | training.logger: wandb |
| Reproducibility Seed | training.random_state: 42 + training.deterministic: true |
| Warmup scheduler | training.scheduler.name: warmup_cosine + warmup_epochs: 10 |
| Config inheritance | _base_: ../base.yaml |
| Albumentations | training.augmentation: albumentations |
| YAML Aug Pipeline | training.augmentation: pipeline + training.aug_pipeline: [...] |
Full config guide: docs/paradigms/README.md
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):
| Category | Methods |
|---|---|
| Geometric | horizontal_flip, vertical_flip, random_rotate90, random_rotate, random_affine, random_perspective, random_scale, elastic_deform, grid_mask |
| Pixel-level | photometric_distortion, color_jitter, brightness_contrast, gamma_correction, clahe, gaussian_blur, gaussian_noise, sharpness, posterize, random_solarize, channel_dropout |
| Masking | random_erasing, coarse_dropout, grid_mask |
| Sample-level | copy_paste, mosaic |
Note: All intensity parameters use
_rangesuffix (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
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
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
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
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
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):
Details: docs/paradigms/text_guided.md
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.
Detailed docs: docs/data/README.md Dataset example configs: configs/intro_to_datasets/
| Type | Description |
|---|---|
synapse | Synapse multi-organ CT (TransUNet format) |
acdc | ACDC cardiac MRI (TransUNet format) |
generic | Generic images/ + masks/ directories |
qata_covid19 | QaTa-COV19 chest X-ray + per-image text (LViT format) |
mosmed_plus | MosMedData+ COVID CT + per-image text (LViT format) |
# 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
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
# 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: {}
# child.yaml — only write overrides
_base_: ../base_resnet50.yaml
model:
num_classes: 9
training:
epochs: 300
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
# 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
@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
@LOSS_REGISTRY.register("my_loss")
class MyLoss(nn.Module):
def forward(self, pred, target):
return loss_value
# 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.
@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},
}
Apache 2.0. For legitimate academic research and engineering use only. Clinical deployment must comply with local regulations.
Thanks to PyTorch, timm, MONAI, SSL4MIS, SAM, GroundingDINO, DINOv2, CLIP, transformers, and all open-source projects that made this possible.
For questions, collaborations, or bug reports, feel free to reach out:
![]() Scan to join QQ Group |
![]() Join our Slack workspace |
Python
99.3%
A Modern Modular 2D Medical Image Segmentation Toolbox
52
stars
45
commits
Python
primary language
Aug 14, 2026
updated
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
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.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).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 .
# 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
Three levels of weight loading, from lightest to heaviest:
| Method | YAML Key | Scope | Use Case |
|---|---|---|---|
| Encoder pretrained | encoder.pretrained: true | Backbone only | ImageNet / domain-specific backbone weights |
| Manual pretrained path | encoder.pretrained_path: /path/to/weights.pth | Backbone only | Offline / custom backbone checkpoint |
| Transfer learning | model.transfer_learning_path: /path/to/full_model.pth | Entire network | Load 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 whenpretrained: false. All auto-downloadable weights support manual path override viapretrained_pathortransfer_learning_path.
| Category | Mechanism | Models |
|---|---|---|
| A. WEIGHT_REGISTRY auto-download | ensure_weight() with GitHub/GCS/HF sources | swinunet, h2former, hiformer, transunet, vm_unet, rwkv_unet (B/S/T), cswin_unet, da_transunet, mamba_unet, fcbformer, transnuseg |
| B. timm / torchvision runtime | pretrained: true triggers built-in download | segformer_b0–b5, esfpnet, cascade, emcad, polyp_pvt, fatnet, transfuse, mist, hsnet, ssformer, ldnet, dconnnet, cfanet, lv_unet, nulite, polyper |
| C. SAM family | pretrained: true auto-downloads ViT/SAM weights | sam_b, sam_l, mobile_sam, sam2, sam_med2d, samed, sammed2d_wrapper, samus, auto_sam, lite_medsam, medical_sam_adapter |
| D. Foundation & MLLM encoders | HuggingFace Hub / open_clip / transformers auto-download | 39 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) |
# 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:
pretrained_path in YAML to point to a local filetimm encoder weights are downloaded automatically via timm's built-in mechanism.
# 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
# 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
# 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
# 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
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
# 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
# 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
# FLOPs / Params / FPS
python profile_model.py --config configs/architectures/networks/general/transunet.yaml
python scripts/export_onnx.py \
--config configs/architectures/networks/general/transunet.yaml \
--checkpoint output/best_model.pth \
--output model.onnx --verify
python scripts/visualize.py \
--config configs/architectures/networks/general/transunet.yaml \
--checkpoint output/best_model.pth \
--input ./data/test/images/ \
--output vis_output/
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")
A step-by-step tutorial series covering deep learning medical image segmentation from fundamentals to advanced topics:
| Chapter | Title | Key Topics |
|---|---|---|
| 01 | Introduction to Medical Image Segmentation | Concepts, clinical significance, metrics, method evolution |
| 02 | U-Net in Detail | Architecture, skip connections, U-Net family variants |
| 03 | Data and Preprocessing | Formats, split strategies, augmentation pipeline |
| 04 | Training and Evaluation | Loss functions, optimizers, AMP/DDP, evaluation |
| 05 | Encoder Deep Dive | CNN / Transformer / Mamba / RWKV comparison, timm wrapper |
| 06 | Decoders and Skip Connections | CASCADE / EMCAD / Attention Gate, skip taxonomy |
| 07 | Foundation Models | DPT head, 9 medical modalities, fine-tuning strategies |
| 08 | Advanced Training Paradigms | Overview: semi-supervised, domain adaptation, distillation, weakly supervised |
| 08a | Semi-Supervised Segmentation | Mean Teacher, CPS, UniMatch, FixMatch, consistency regularization |
| 08b | Domain Adaptation | AdvEnt, DANN, TENT, FDA, MIC, HRDA, SePiCo |
| 08c | Knowledge Distillation | VanillaKD, DKD, CWD, MGD, DIST, ReviewKD, SimKD |
| 08d | Weakly Supervised Segmentation | CAM, SEAM, PuzzleCAM, Box, Point, Scribble supervision |
| 08e | Text-Guided Segmentation | CRIS, BiomedParse, LViT, CLIP-based, MLLM pipeline |
| 09 | Deployment and Inference | ONNX export, TTA, ensemble, MLLM pipeline |
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
Detailed docs: docs/models/
| Category | Count | Examples |
|---|---|---|
| CNN | 36 | UNet, UNet3+, UNet++, Attention-UNet, nnU-Net, MedNeXt, MEW-UNet, DCSAU-Net |
| Transformer | 36 | SegFormer, TransUNet, Swin-UNet, DAEFormer, MISSFormer, HiFormer, PolypPVT, CASCADE |
| Mamba / SSM | 24 | VM-UNet, U-Mamba, Swin-UMamba, LKM-UNet, LoG-VMamba, HC-Mamba |
| SAM family | 10 | MedSAM, SAM-Med2D, SAM2, SAMUS, AutoSAM, MobileSAM, LiteMedSAM, SAMed, Medical SAM Adapter |
| KAN / MLP | 4 | RollingUNet, UNeXt, U-KAN, WA-UKAN |
| Linear Attention | 3 | TTT-UNet, U-VixLSTM, xLSTM-UNet |
| RWKV | 5 | U-RWKV (MICCAI 2025), U-RWKV (TIP 2026), RWKV-UNet, MD-RWKV-UNet, RIR-Zigzag |
| Text-guided | 12 | CRIS, 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_rwkv— MICCAI 2025: Direction-Adaptive RWKV Module (DARM) + Stage-Adaptive Squeeze-and-Excitation (SASE), lightweight design with RWKV integrated within conv stages. Source: hbyecoding/U-RWKVu_rwkv_tip— IEEE TIP 2026: Standard U-Net + post-conv RWKV attention blocks with OmniShift multi-scale conv, originally for volumetric segmentation. Source: hbyecoding/U-RWKV
Highlight: 39 foundation model encoders covering 9 medical modalities
| Modality | Count | Models |
|---|---|---|
| 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, RETFound, FLAIR, OphMAE |
| Dermatology | 3 | DermCLIP, MoNet, PanDerm |
| General Medical | 3 | BiomedCLIP, MedCLIP, MedSigLIP |
| MLLM Vision | 8 | Qwen2.5-VL, Qwen3-VL, MedGemma, LLaVA-Med, HuatuoGPT, HealthGPT, HuLuMed, LingShu |
| Ultrasound | 3 | UltraFedFM, USF-MAE, SAMUS |
| Endoscopy | 3 | EndoViT, 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
| Category | Count | Examples |
|---|---|---|
| Basic (upsampling) | 4 | deconv_upcat (unet), Bilinear, deconv_catup (deconv), DepthwiseSep |
| Dense (connections) | 2 | UNet++, UNet3+ |
| Cascade | 10 | CASCADE, EMCAD (2 variants), G-CASCADE (2 variants), CFM, MERIT (2 variants), EDLDNet |
| Attention | 6 | Attention Gate, BANet, CCNet, Lawin, OCRNet, UCTransNet |
| Transformer | 5 | DAEFormer, MISSFormer, MTUNet, SwinUNet, nnFormer |
| MLP | 2 | SegFormer MLP, MLP Decoder |
| Specific (network) | 13 | CFA-Net, DCSA-UNet, EGNet, FAT-Net, FF-Parser, H2Former, HAM, HiFormer, KI-UNet, MALUNet, RWKV-UNet, ScaleFormer, TransUNet |
| Mamba | 1 | VM-UNet |
| Pyramid | 2 | UPerNet, DeepLabV3 (ASPP) |
Full list: docs/models/decoders.md
Detailed docs: docs/paradigms/
| Feature | YAML config |
|---|---|
| Mixed precision AMP | training.amp: true or CLI --amp |
| Multi-GPU DDP | torchrun --nproc_per_node=N train.py |
| DataParallel | training.parallel: dp |
| TensorBoard | training.logger: tensorboard |
| WandB | training.logger: wandb |
| Reproducibility Seed | training.random_state: 42 + training.deterministic: true |
| Warmup scheduler | training.scheduler.name: warmup_cosine + warmup_epochs: 10 |
| Config inheritance | _base_: ../base.yaml |
| Albumentations | training.augmentation: albumentations |
| YAML Aug Pipeline | training.augmentation: pipeline + training.aug_pipeline: [...] |
Full config guide: docs/paradigms/README.md
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):
| Category | Methods |
|---|---|
| Geometric | horizontal_flip, vertical_flip, random_rotate90, random_rotate, random_affine, random_perspective, random_scale, elastic_deform, grid_mask |
| Pixel-level | photometric_distortion, color_jitter, brightness_contrast, gamma_correction, clahe, gaussian_blur, gaussian_noise, sharpness, posterize, random_solarize, channel_dropout |
| Masking | random_erasing, coarse_dropout, grid_mask |
| Sample-level | copy_paste, mosaic |
Note: All intensity parameters use
_rangesuffix (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
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
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
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
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
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):
Details: docs/paradigms/text_guided.md
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.
Detailed docs: docs/data/README.md Dataset example configs: configs/intro_to_datasets/
| Type | Description |
|---|---|
synapse | Synapse multi-organ CT (TransUNet format) |
acdc | ACDC cardiac MRI (TransUNet format) |
generic | Generic images/ + masks/ directories |
qata_covid19 | QaTa-COV19 chest X-ray + per-image text (LViT format) |
mosmed_plus | MosMedData+ COVID CT + per-image text (LViT format) |
# 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
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
# 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: {}
# child.yaml — only write overrides
_base_: ../base_resnet50.yaml
model:
num_classes: 9
training:
epochs: 300
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
# 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
@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
@LOSS_REGISTRY.register("my_loss")
class MyLoss(nn.Module):
def forward(self, pred, target):
return loss_value
# 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.
@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},
}
Apache 2.0. For legitimate academic research and engineering use only. Clinical deployment must comply with local regulations.
Thanks to PyTorch, timm, MONAI, SSL4MIS, SAM, GroundingDINO, DINOv2, CLIP, transformers, and all open-source projects that made this possible.
For questions, collaborations, or bug reports, feel free to reach out:
![]() Scan to join QQ Group |
![]() Join our Slack workspace |
Python
99.3%