ComWin

This repository contains data preparation, training and evaluation code, and pretrained models for "Compete to Win: Enhancing Pseudo Labels for Barely-supervised Medical Image Segmentation".

Introduction

This study investigates barely-supervised medical image segmentation where only few labeled data, i.e., single-digit cases are available. We observe the key limitation of the existing state-of-the-art semi-supervised solution cross pseudo supervision is the unsatisfactory precision of foreground classes, leading to a degenerated result under barely-supervised learning. In this paper, we propose a novel Compete-to-Win method (ComWin) to enhance the pseudo label quality. In contrast to directly using one model’s predictions as pseudo labels, our key idea is that high-quality pseudo labels should be generated by comparing multiple confidence maps produced by different networks to select the most confident one (a compete-to-win strategy). To further refine pseudo labels at near-boundary areas, an enhanced version of ComWin, namely, ComWin+, is proposed by integrating a boundary-aware enhancement module. Experiments show that our method can achieve the best performance on three public medical image datasets for cardiac structure segmentation, pancreas segmentation and colon tumor segmentation, respectively.

Usage

##Requirements

pip install -r requirements.txt

Data preparation

Download Pancreas dataset Version 2 and prepare h5 files and data splits by conducting:

raw_data_dir="where you put Pancreas-CT data"
raw_label_dir="where you put Pancreas-CT labels"
cd code/dataloaders
python pancreas_processing.py --raw_data_dir $raw_data_dir --raw_label_dir $raw_label_dir

Training

cps_un_weight_final=0.5 # or other numbers
window_size=4 # or other numbers
python train_comwin_plus.py --gpu 0 --labeled_num=3 --window_size=$window_size --sparse_attn=True --cps_un_weight_final=$cps_un_weight_final --cps_la_weight_final=0 --exp=try_comwin_plus --root_path ../data/Pancreas-CT-all

Test with checkpoints you trained

model_name=try_comwin_plus # or other names
python test_comwin_plus.py --root_path ../data/Pancreas-CT-all --image_list_path=pancreas_test.list --ds_starting_layer=8 --window_size=$window_size --sparse_attn=True --gpu 0 --iter 6000 --model=$model_name

Acknowledgement

We would like to thank following open-source projects: UA-MT, CPS, and SwinTransformer.