Modelpublic

MrCzaro/Pressure_sore_cascade_classifier_Torch

sourceHugging Facemitupdated 7mo agoView on Hugging Face

0likes

Model Card

Pressure Sore Cascade Classifier — Torchvision 3-Level

This repository contains 8 PyTorch model weights forming a 3-level hierarchical cascade for automated pressure sore detection and staging. The cascade progressively narrows from detecting any wound, to separating severity groups, to fine-grained stage classification — mirroring clinical decision-making.

These weights are used by ps_classifier_torch_cascade.py in the PS Classifier web application.

Cascade Structure

Image
  │
  ▼
[Level 1 — PS vs No-PS]                         BCEWithLogitsLoss · sigmoid
  MaxVit_T  (linear head)
  ResNet50  (mlp head)
  │
  ├─ NO  → "No pressure sore detected"
  └─ YES ▼
[Level 2 — Early (Stage I/II) vs Advanced (III/IV)]   BCEWithLogitsLoss · sigmoid
  ConvNeXt_Base      (mlp head)
  EfficientNet_V2_L  (linear head)
  │
  ├─ EARLY ──────────────────────┐
  └─ ADVANCED ───────┐           │
                     ▼           ▼
  [Level 3b]                   [Level 3a]
  Stage III vs Stage IV         Stage I vs Stage II
  ConvNeXt_Large (MSH)          EfficientNet_V2_L (mlp)
  ViT_B_16 (mlp)                ConvNeXt_Tiny (linear)
  CrossEntropyLoss              BCEWithLogitsLoss · sigmoid
  WrappedModel pattern          Direct-attachment pattern
  ↓ Confidence gate 0.65 ↓      ↓ Confidence gate 0.65 ↓

Confidence gating: if the Level 3 ensemble confidence falls below 0.65 the prediction is still returned, but annotated with an uncertainty warning and flagged for clinical review (details["level_3"]["gated"] == True).

Files

File	Level	Architecture	Head	Loss
`Level 1 Binary PS or not PS MaxVit_T.pth`	L1	MaxVit-T	linear	BCE
`Level 1 Binary PS or not PS ResNet50.pth`	L1	ResNet-50	mlp	BCE
`Level 2 Early vs Advanced ConvNeXt_Base.pth`	L2	ConvNeXt-Base	mlp	BCE
`Level 2 Early vs Advanced EfficientNet_V2_L.pth`	L2	EfficientNet-V2-L	linear	BCE
`Level 3a Early EfficientNet_V2_L.pth`	L3a	EfficientNet-V2-L	mlp	BCE
`Level 3a Early ConvNeXt_Tiny.pth`	L3a	ConvNeXt-Tiny	linear	BCE
`Level 3b Advanced ConvNeXt_Large.pth`	L3b	ConvNeXt-Large	multistagehead	XEnt
`Level 3b Advanced ViT_B_16.pth`	L3b	ViT-B/16	mlp	XEnt

MSH = MultiStageHead (Dropout → FC(in→in/2) → BN → ReLU → FC(in/2→2))

Model Performance

Level 1 — PS vs No-PS (test set: 261 images)

Model	Head	Dropout	Scheduler	Accuracy	Macro F1	AUC-ROC
MaxVit_T	linear	0.2396	CosineAnnealingLR	0.9962	0.9962	1.0000
ResNet50	mlp	0.5840	CosineAnnealingLR	1.0000	1.0000	1.0000

Level 2 — Early vs Advanced (test set: 125 images)

Model	Head	Dropout	Optimizer	Scheduler	Accuracy	Macro F1	AUC-ROC
ConvNeXt_Base	mlp	0.1025	AdamP	CosineAnnealingLR	0.9520	0.9520	0.9857
EfficientNetV2L	linear	0.3564	AdamP	CosineAnnealingLR	0.9600	0.9600	0.9916

Level 3a — Stage I vs Stage II (test set: 63 images)

Model	Head	Dropout	Optimizer	Scheduler	Accuracy	Macro F1	AUC-ROC
EfficientNetV2L	mlp	0.1949	AdamP	StepLR	0.9048	0.9047	0.9849
ConvNeXt_Tiny	linear	0.1601	Lion	CosineAnnealingLR	0.9683	0.9682	0.9909

Level 3b — Stage III vs Stage IV (test set: 63 images)

Model	Head	Dropout	Optimizer	Scheduler	Accuracy	Macro F1	AUC-ROC
ConvNeXt_Large	multistagehead	0.6594	Lion	ReduceLROnPlateau	0.7778	0.7773	0.8861
ViTB16	mlp	0.5445	AdamW	ReduceLROnPlateau	0.7937	0.7934	0.8569

Stage III vs Stage IV is the hardest sub-task — subtle visual differences between full-thickness tissue loss with and without exposed bone/muscle make it challenging even for clinicians. The confidence gate at Level 3b (0.65) flags the most uncertain predictions.

Training Details

All models were trained on a curated dataset of ~1,000 pressure sore images collected from public medical databases, with stratified 70/20/10 train/validation/test splits.

Shared configuration:

—Input: 224 × 224, ImageNet normalisation (mean [0.485, 0.456, 0.406], std [0.229, 0.224, 0.225])
—Augmentation: random flips, rotation, colour jitter, Gaussian blur, affine transforms (Albumentations)
—Freeze schedule: backbone frozen for initial epochs, then progressively unfrozen (2-stage)
—Early stopping: patience 8–10 epochs on validation loss
—Hyperparameters: selected by Optuna trials (learning rate, weight decay, dropout, head type)
—Mixed precision: fp16 via Accelerate

Architecture notes:

—L1, L2, L3a: head is attached directly to the backbone's native classifier slot (model.classifier[2] for ConvNeXt, model.heads.head for ViT, etc.). Saved state dict has flat keys.
—L3b: WrappedModel wrapper — backbone classifier slot replaced with nn.Identity, a separate head receives raw feature embeddings. Saved state dict has backbone.* / head.* key prefixes.

Usage

Installation

bash

pip install torch torchvision albumentations pillow

Minimal inference example

python

import torch
import torch.nn as nn
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torchvision import models
from PIL import Image

# Helpers 

def load_standard(arch_fn, in_feat, head_type, dropout, path, num_classes=1):
    """L1 / L2 / L3a — head directly on backbone, BCE/sigmoid."""
    model = arch_fn(weights=None)
    if head_type == "linear":
        head = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_feat, num_classes))
    else:  # mlp
        head = nn.Sequential(
            nn.Dropout(dropout), nn.Linear(in_feat, in_feat // 2),
            nn.ReLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(in_feat // 2, num_classes))
    model.classifier[2] = head   # adjust slot per arch (see ps_classifier_torch_cascade.py)
    sd = torch.load(path, map_location="cpu", weights_only=False)
    model.load_state_dict(sd, strict=False)
    return model.eval()


transform = A.Compose([
    A.Resize(224, 224),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

def preprocess(path):
    img = Image.open(path).convert("RGB")
    return transform(image=np.array(img))["image"].unsqueeze(0)

# Full cascade (recommended: use ps_classifier_torch_cascade.py) 

# Clone the repo and place weights under models/torch_cascade/
# then simply:

from ps_classifier_torch_cascade import classify_image_cascade, cascade_confidence

image, message, details = classify_image_cascade("path/to/wound.jpg")
print(message)
# ✅ Pressure sore detected
# Severity : early (0.96)
# Stage    : Stage II (0.91)

print("Joint confidence:", cascade_confidence(details))  # e.g. 0.864

if details.get("level_3", {}).get("gated"):
    print("⚠ Low Level-3 confidence — recommend clinical review")

`details` schema