PromptMask

A Controllable Text-Driven Masking Framework
for Self-Supervised Visual Pretraining

ICME 2026 · Official PyTorch implementation

Paper (ICME) · arXiv · Results · Quick start · Documentation · Citation

Tell the model what matters. Let masking focus the learning.

We introduce PromptMask, a text-guided masked image modeling framework for medical imaging. We localize task-relevant regions with a vision-language model, then mask these regions more heavily than the background. This brings task semantics into self-supervised pretraining without requiring per-image reports or annotations.

Text prompts guide region localization; region-specific masking guides representation learning.

Text-guided control	Transfer across modalities	Lower masking ratios
Guide localization with category phrases or visual descriptions.	Evaluate on brain MRI, chest CT, and lung X-ray; transfer to classification, detection, and instance segmentation.	On chest CT, we reach 96.8% accuracy at 40% masking, compared with SparK's 94.4% at 70%.

Results

We report our original experiments from the ICME manuscript below. All values, baselines, and ablations are available in CSV and JSON, with their source table or figure. Accuracy and AP are shown on a 0–100 scale.

Classification

Method	Brain MRI Fine-tuning	Chest CT Fine-tuning	Lung X-ray Fine-tuning	Brain MRI Linear probe	Chest CT Linear probe	Lung X-ray Linear probe
SparK	96.2	94.4	94.7	69.8	63.5	81.6
AnatoMask	96.5	95.8	95.2	70.2	63.9	85.3
PromptMask	96.8	97.5	96.0	71.3	67.0	88.5

ICME Tables II–III. We show the closest convolutional masked-modeling baselines here; the complete fine-tuning table also includes MoCoV2, BYOL, SimCLR, and MAE.

Detection and instance segmentation

Method	Box AP	Box AP75	Mask AP	Mask AP75
SparK	45.6	53.2	44.8	52.0
AnatoMask	45.2	53.1	46.2	54.1
PromptMask	46.9	53.5	45.9	54.1

ICME Table IV, BR35H. Our gains over SparK are +1.3 Box AP and +1.1 Mask AP.

Masking efficiency and qualitative examples

At 40% masking, we exceed SparK at 70% masking. The plot reports classification accuracy; it does not measure wall-clock speed.

See all numerical results and figure descriptions.

How it works

1. Localize with text. We encode images and prompts with BiomedCLIP, obtain saliency with M2IB, and extract connected foreground regions. SAM refines the localization before bounding-box expansion.
2. Mask by region. We assign separate masking ratios to the region of interest and the background.
3. Learn and transfer. A sparse encoder and hierarchical decoder reconstruct masked image patches. We transfer the encoder to classification, detection, and instance segmentation.

Our implementation includes ResNet-50 and a ViT-S adapter, local model loading, offline prompt preparation, checkpoint recovery, and configurable ablation runs. Method and implementation details →

Quick start

From the repository root, use Python 3.11 or newer:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e '.[ml,data,detection]'

# Check the configuration without loading a model.
python scripts/check_config.py

# Run the full localization → masking → reconstruction interface on CPU fixtures.
python scripts/smoke_m6.py

The smoke test uses two synthetic images and lightweight model fixtures. It needs no dataset or model download and returns M6_SMOKE_PASS. For a CPU-only PyTorch installation or optional model packages, see Getting started.

Run experiments

We separate data preparation, localization, pretraining, and downstream evaluation into explicit commands:

Step	Entry point	Guide
Prepare local datasets	scripts/prepare_data_manifest.py	Data sources · Manifest format
Prepare prompts and regions	scripts/generate_prompts.py, scripts/localize.py	Model assets · Experiments
Pretrain	scripts/pretrain.py	Pretraining and resume
Fine-tune, probe, detect, or segment	scripts/downstream.py	Downstream experiments
Run ablations	scripts/ablations.py	Ablation configuration

Use --help for each entry point. We keep machine-specific paths and experiment choices in local TOML overlays. Training commands first describe the execution plan; --execute starts the configured run. Dataset and model sources are documented, while their assets remain outside the repository.

Documentation

• Getting started: installation and CPU smoke tests.
• Experiment guide: localization, pretraining, transfer, evaluation, and resume.
• Data sources and formats: dataset access and input preparation.
• Configuration and model assets: overrides, local checkpoints, and adapters.
• Paper results: all ICME tables and the masking-ratio curve in reusable formats.

Repository structure

promptmask/
├── localization/   # Text embeddings, M2IB, region extraction, refinement
├── mim/            # Region masking, reconstruction, pretraining
├── models/         # VLM, SAM, sparse CNN, ViT, checkpoint adapters
├── downstream/     # Classification, detection, segmentation, COCO AP
├── data/           # Medical image readers and dataset manifests
└── experiments/    # Experiment runners, caching, checkpoints, ablations
configs/            # Base settings, task overlays, and smoke fixtures
scripts/            # Data preparation, experiment, and test entry points
results/paper/      # Numerical results from our ICME manuscript
assets/figures/     # Framework and result figures

Tests

python -B scripts/validate_all.py

We provide CPU regression tests and seven smoke checks for the method and experiment interfaces. The command records results in docs/VALIDATION.json and docs/VALIDATION.log; generated reports and experiment outputs are ignored by Git. A GitHub Actions workflow runs the same checks.

Citation

Our public preprint uses the earlier title Mask What Matters. Please cite its published record:

@misc{wang2025maskmatters,
  title  = {Mask What Matters: Controllable Text-Guided Masking for Self-Supervised Medical Image Analysis},
  author = {Ruilang Wang and Shuotong Xu and Bowen Liu and Runlin Huang and Donglong Chen and Weifeng Su},
  year   = {2025},
  eprint = {2509.23054},
  archivePrefix = {arXiv},
  primaryClass = {cs.CV},
  url    = {https://arxiv.org/abs/2509.23054}
}

We use the ICME manuscript for the method and results in this repository. Paper versions →

Acknowledgements and license

We build on BiomedCLIP, M2IB, SAM, and SparK, and use the PyTorch ecosystem. We thank the researchers and maintainers who make these resources available. Third-party references →

Our code is released under the MIT license. External datasets, model assets, and manuscript content retain their respective terms.