Introduction

This repository is the official implementation of our paper: PLESS: Pseudo-Label Enhancement with Spreading Scribbles for Weakly Supervised Segmentation.

Preprint: arXiv:2602.11628 | Journal publication: (coming soon)

---

Citation

If you use this work, please cite:

@misc{gabrielyan2026pless,
  title         = {PLESS: Pseudo-Label Enhancement with Spreading Scribbles for Weakly Supervised Segmentation},
  author        = {Yeva Gabrielyan and Varduhi Yeghiazaryan and Irina Voiculescu},
  year          = {2026},
  eprint        = {2602.11628},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  url           = {https://arxiv.org/abs/2602.11628}
}

---

Dataset

ACDC (Automated Cardiac Diagnosis Challenge)

The project uses the ACDC cardiac MRI dataset with scribble annotations for weakly-supervised segmentation.

Obtaining the data:

1. Full mask annotations: ACDC Challenge
2. Scribble annotations: MAAG scribbles
3. Pre-processed slices and volumes (HDF5 format): WSL4MIS / ScribbleVC or CycleMix

Expected directory layout:

data/
└── ACDC/
    ├── ACDC_training_slices/   # 2D HDF5 slices for training
    ├── ACDC_training_volumes/  # 3D HDF5 volumes for validation during training
    ├── ACDC_testing/           # 3D HDF5 volumes for final evaluation
    └── slice_classification.xlsx

Place the data/ folder one level above the code/ directory (i.e., alongside it).

Dataset details:

• 4 classes: background (0), RV (1), MYO (2), LV (3)
• Each HDF5 file contains: image, label (full mask), scribble (weak annotation)

MSCMR (Multi-sequence Cardiac MRI)

1. Dataset and scribble annotations: MSCMRseg and CycleMix
2. Pre-processed files: MSCMR dataset

Expected directory layout:

data/
└── MSCMR/
    ├── MSCMR_training_slices/     # 2D HDF5 slices for training
    ├── MSCMR_validation_volumes/  # 3D HDF5 volumes for validation during training
    ├── MSCMR_testing/             # 3D HDF5 volumes for final evaluation
    └── slice_classification.xlsx

---

Requirements

Training and evaluation:

Python >= 3.8
torch
torchvision
numpy
scipy
h5py
nibabel
SimpleITK
medpy
scikit-image
opencv-python
tensorboardX
pandas
tqdm

pip install numpy scipy h5py nibabel SimpleITK medpy scikit-image opencv-python tensorboardX pandas tqdm

Preprocessing (PLESS enhancement only):

pip install pycuda

PyCUDA requires a working CUDA installation. The waterfall kernels are compiled at runtime via PyCUDA — no separate build step is needed.

---

Preprocessing — PLESS Enhancement

PLESS improves pseudo-labels by spreading scribble annotations across watershed regions. Before training with --PLESS 1, you must generate the enhanced scribble H5 files from the original training slices.

Run from the code/PLESS_preprocessing/ directory:

python preprocess.py ../../data/ACDC/ACDC_training_slices ../../data/ACDC/ACDC_training_slices_ENH

For MSCMR:

python preprocess.py ../../data/MSCMR/MSCMR_training_slices ../../data/MSCMR/MSCMR_training_slices_ENH

Usage: python preprocess.py <input_dir> <output_dir>

The script skips files already present in <output_dir>, so it is safe to resume after interruption.

After preprocessing, the data layout should be:

data/
└── ACDC/
    ├── ACDC_training_slices/      # original slices
    ├── ACDC_training_slices_ENH/  # enhanced scribble slices (generated above)
    ├── ACDC_training_volumes/
    ├── ACDC_testing/
    └── slice_classification.xlsx

---

Training

Run training from the code/ directory:

python train.py --gpu 0 --exp model1

Key arguments:

Argument	Default	Description
--root_path	../data/ACDC	Path to the dataset root
--data	ACDC	Dataset to use (ACDC or MSCMR)
--exp	ACDC/debug	Experiment name (used for saving checkpoints and logs)
--fold	fold1	Data split (fold1–fold5)
--sup_type	scribble	Supervision type (scribble or label)
--tau	0.5	Confidence threshold for pseudo-label generation
--max_iterations	60000	Total training iterations
--batch_size	12	Batch size
--base_lr	0.01	Initial learning rate
--gpu	0	GPU device ID
--seed	2022	Random seed
--PLESS	0	Enable PLESS pseudo-label enhancement (1) or disable (0)
--enh_train_dir	None	Subdirectory of enhanced H5 slices; required when --PLESS 1
--tolerance	0.5	Fraction of total epochs during which PLESS enhancement is applied

Standard training (scribble supervision only):

python train.py \
  --root_path ../data/ACDC \
  --exp baseline/run1 \
  --fold fold1 \
  --gpu 0

PLESS training (with waterfall-enhanced scribbles):

python train.py \
  --root_path ../data/ACDC \
  --exp pless/run1 \
  --fold fold1 \
  --PLESS 1 \
  --enh_train_dir /ACDC_training_slices_ENH \
  --tolerance 0.25 \
  --tau 0.5 \
  --max_iterations 60000 \
  --batch_size 12 \
  --gpu 0

Checkpoints and TensorBoard logs are saved under ../model/{exp}_{fold}/{sup_type}/.

---

Evaluation

Run 3D evaluation from the code/ directory:

python test.py \
  --test_dir ../data/ACDC/ACDC_testing \
  --exp pless/run1 \
  --fold fold1 \
  --gpu 0

The script loads every .h5 volume from --test_dir, runs inference, and reports per-class metrics. The model checkpoint is loaded from ../model/{exp}_{fold}/{sup_type}/unet_best_model.pth.

Key arguments:

Argument	Default	Description
--test_dir	required	Directory of test H5 volumes — all .h5 files are evaluated
--exp	ACDC/debug	Experiment name (must match the training run)
--fold	fold1	Fold used during training (needed to locate the checkpoint)
--sup_type	scribble	Supervision type (must match the training run)
--ori_img_path	None	Directory with .nii.gz files for voxel spacing; defaults to (1,1,1) if omitted
--num_classes	4	Number of output classes
--gpu	0	GPU device ID

Metrics reported per class (RV, MYO, LV):

• Dice Coefficient (DC)
• Hausdorff Distance 95th percentile (HD95)
• Average Surface Distance (ASD)

Predictions are saved to ../model/{exp}_{fold}/{sup_type}/unet_predictions/.

---

Acknowledgements

Our code is based on ScribbleVS, which itself builds upon ScribbleVC and CycleMix. Thanks to their authors for the valuable open-source contributions.