--- license: mit library_name: pytorch pipeline_tag: image-classification tags: - medical-imaging - wound-care - wound-segmentation - vision-transformer - dinov2 - convnext - sam2 - multi-task - mlhc-2026 --- # WILLIE — Wound Classification, Segmentation & Localization
WILLIE architecture overview 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](https://github.com//) - **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 ```python 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](https://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](http://www.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 ```bibtex @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.