Debiasing Medical Knowledge for Prompting Universal Model in CT Image Segmentation (TMI)

by Boxiang Yun, Shitian Zhao, Qingli Li, Alex Kot, Yan Wang*

Introduction

A PyTorch implementation of DUM, a causal debiasing framework that enhances generalization in universal medical image segmentation by mitigating knowledge bias from text prompts. DUM leverages both organ-level semantic priors and instance-level visual context to improve robustness across diverse clinical scenarios. Based on the paper: Debiasing Medical Knowledge for Prompting Universal Model in CT Image Segmentation.

Requirements

This repository is based on PyTorch 2.2.2, CUDA 11.8, and Python 3.9.7. All experiments in our paper were conducted on NVIDIA GeForce RTX 3090 GPU with an identical experimental setting.

Data Preparation

• 01 Multi-Atlas Labeling Beyond the Cranial Vault - Workshop and Challenge (BTCV)
• 02 Pancreas-CT TCIA ## The label we used for Dataset 01 and 02 is in here.
• 03 Combined Healthy Abdominal Organ Segmentation (CHAOS)
• 04 Liver Tumor Segmentation Challenge (LiTS)
• 05 Kidney and Kidney Tumor Segmentation (KiTS)
• 06 Liver segmentation (3D-IRCADb)
• 07 WORD: A large scale dataset, benchmark and clinical applicable study for abdominal organ segmentation from CT image
• 08 AbdomenCT-1K
• 09 Multi-Modality Abdominal Multi-Organ Segmentation Challenge (AMOS)
• 10 Decathlon (Liver, Lung, Pancreas, HepaticVessel, Spleen, Colon
• 11 CT volumes with multiple organ segmentations (CT-ORG)
• 12 AbdomenCT 12organ

The post_label can be downloaded via link. (see https://github.com/ljwztc/CLIP-Driven-Universal-Model/blob/main/README.md)

• Directory Structure

Datasets/
├── 01_Multi-Atlas_Labeling
│   ├── img
│   │   ├── img0001.nii.gz
│   │   └── ...
│   ├── label
│   └── post_label
├── 02_TCIA_Pancreas-CT
├── ...
├── 15_3DIRCADb
├── 16_TotalsegmentatorV2
└── 17_AbdomenAtlas_1_0_Mini

Usage

Train:

To train a model,

CUDA_VISIBLE_DEVICES=0,1,2,3 python -W ignore -m torch.distributed.launch --nproc_per_node=4 --master_port=12340 \
train_main_dum.py --dist True --num_workers 8 --num_samples 2 --uniform_sample \
  --log_name train_base_4gpu_1002 --data_root_path ../Datasets/ \
  --num_context 4 --alpha 0.5 --batch_size 1 --max_epoch 1000

• Checkpoints and TensorBoard saved in ./out/train_base_4gpu_1002/

To test a model,

CUDA_VISIBLE_DEVICES=0 python -W ignore test.py \
--resume1 ./out/train_base_4gpu_1002/epoch_x.pth \
--data_root_path ../Datasets/ \
--dataset_list PAOT_10_inner --num_workers 8

Citation

If you find the project useful, please consider citing:

@article{yun2025debiasing,
  title={Debiasing Medical Knowledge for Prompting Universal Model in CT Image Segmentation},
  author={Yun, Boxiang and Zhao, Shitian and Li, Qingli and Kot, Alex and Wang, Yan},
  journal={IEEE Transactions on Medical Imaging},
  year={2025},
  publisher={IEEE}
}

TODO List

• Release pretrained model weights and provide a simple inference pipeline.
• Release well-organized and preprocessed 12 datasets to easy adoption