WILLIE — Wound Classification, Segmentation & Localization
One framework, three tasks: classification, segmentation, and segmentation-derived localization over a shared 5-class wound taxonomy.
Checkpoints for WILLIE, a unified vision-transformer framework for wound classification, segmentation and localization, published at MLHC 2026.
Developed in the Qian Group, University of Houston.
- Code, manifests and reproduction notebooks: GitHub repository
- Paper: MLHC 2026 (add link when available)
Repository contents
mini/ willie_mini_fold0-4_best.pt 5-fold, 34.3M params
base/ willie_base_fold0-4_best.pt 5-fold, 520.4M params
xl/ willie_xl_fold0-4_best.pt 5-fold, 762.5M params
decoders/ medsam_decoder_best.pt fine-tuned MedSAM mask decoder
sam2_decoder_best.pt fine-tuned SAM2 mask decoder
Each scale ships all five cross-validation folds. The paper's reported numbers use the fold ensemble, not a single fold.
| Scale | Params | Backbones | Input |
|---|---|---|---|
| MINI | 34.3M | DINOv2-S + FPN | 224×224 |
| BASE | 520.4M | + ConvNeXt-L (dual backbone) | 378×378 |
| XL | 762.5M | + SAM2-Hiera-L (triple backbone) | 378×378 |
Architecture components: F²DCA, WA-CSA, MoE-8, WTCS, WBRN. See the paper for definitions and the ablation.
Results
Test-set results as reported in the paper.
| Task | Metric | Score |
|---|---|---|
| Classification | Accuracy | 91.88% |
| Segmentation | Dice | 91.41% |
| Localization | AP@0.5 | 96.23% |
5-fold cross-validation (mean ± std):
| Scale | Cls Acc | Seg Dice | Det AP@0.5 |
|---|---|---|---|
| MINI | 86.4 ± 1.2 | 83.6 ± 1.7 | 85.4 ± 2.7 |
| BASE | 88.5 ± 2.0 | 87.5 ± 1.0 | 85.9 ± 4.5 |
| XL | — | 91.5 ± 1.1 | — |
CV figures are lower than the headline numbers because the headline uses the held-out test split with test-time augmentation and fold ensembling. Both protocols are described in the paper.
Usage
import torch
from huggingface_hub import hf_hub_download
path = hf_hub_download("QianGroup/willie-weights", "xl/willie_xl_fold0_best.pt")
ckpt = torch.load(path, map_location="cpu")
# MINI/BASE checkpoints wrap the weights; XL is a bare state_dict
state = ckpt.get("model_state_dict", ckpt)
model.load_state_dict(state)
model.eval()
Model classes are in the GitHub repository under
willie_app/architectures/ (willie_mini.py, willie_base.py,
willie_xl.py).
Preprocessing. ImageNet normalization
(mean 0.485, 0.456, 0.406, std 0.229, 0.224, 0.225).
Resize 256→224 for MINI, 420→378 for BASE and XL.
Classes (5): diabetic, pressure, surgical, venous, no_wound.
Checkpoint format
MINI and BASE checkpoints contain model_state_dict, epoch, fold,
metrics, config and optimizer_state_dict. XL checkpoints are bare
state dicts. Optimizer state is retained so training can be resumed; drop it
for inference-only use.
Training data
Trained on three public wound datasets, not redistributed here:
| Dataset | Source |
|---|---|
| FUSeg | fusc.grand-challenge.org — requires accepting the challenge data-use agreement |
| AZH | UWM Big Data Lab, AZH Wound and Vascular Center |
| Medetec | medetec.co.uk |
Exact splits are defined by the manifests in the GitHub repository
(3,535 referenced files). A verify_data.py script there confirms placement.
Intended use
Intended for research: reproducing the paper's benchmark, and as a starting point for wound-analysis research.
Not intended for clinical use. This is research code, not a medical device. It has not been validated for diagnosis, treatment decisions or patient care, and carries no regulatory clearance.
Limitations
Duplicate images across splits. Byte-level checks of the source datasets
found one image identical between AZH train and test
(train/surgical/10_0.jpg / test/surgical/99_0.jpg). On the 234-image AZH
test set the maximum effect on the reported 91.88% accuracy is 0.43
percentage points (worst case 91.45%). Nine image pairs are identical between
FUSeg train and validation; validation is used for model selection only and
does not enter the reported test Dice. FUSeg train↔test and validation↔test
contain no duplicates.
Narrow domain. All training images are foot, pressure, venous and surgical wound photographs from a small number of clinical sources. Generalisation to other imaging conditions, skin tones, camera hardware and wound types is untested. Performance on under-represented skin tones has not been measured.
Ablation protocol. Reported component effects fall within cross-validation noise. The XL ablation implementation differs substantially from MINI and BASE, so cross-scale ablation comparisons should be treated with caution.
Fold ensembling. Headline numbers come from the 5-fold ensemble with test-time augmentation. A single fold will score lower.
Licence
Weights released under the MIT Licence, matching the code repository.
The underlying datasets are not covered by this licence and remain subject to their original terms, including the FUSeg challenge data-use agreement. Users are responsible for obtaining the data under its own terms.
Third-party foundation weights (SAM2, MedSAM, DINOv2, ConvNeXt) are not redistributed here — obtain them from their original sources under their own licences.
Citation
@inproceedings{willie2026,
title = {WILLIE: A Unified Vision-Transformer Framework and Benchmark
for Wound Classification, Segmentation, and Localization},
author = {Maddikunta, Gopi Trinadh and Qian, Peizhu},
booktitle = {Proceedings of the Machine Learning for Healthcare Conference (MLHC)},
year = {2026}
}
Please also cite the FUSeg, AZH and Medetec datasets per their own requirements.
Developed in the Qian Group, University of Houston. Advisor: Dr. Peizhu Qian. Computation performed on the UH carya cluster.