Skip to content

Latest commit

 

History

History
80 lines (67 loc) · 4.05 KB

File metadata and controls

80 lines (67 loc) · 4.05 KB

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

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