README.md

---
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

<div align="center">

<img src="assets/WILLIE Workflow.png" alt="WILLIE architecture overview" width="90%">

<em>One framework, three tasks: classification, segmentation, and
segmentation-derived localization over a shared 5-class wound taxonomy.</em>

</div>



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/<GITHUB_ORG>/<REPO>)
- **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.