Add user guidance: pre-upload tips, low-confidence warning, pressure-class skepticism note

app.py

@@ -66,29 +66,56 @@ MODEL = _load_model(CKPT_PATH)
 TRANSFORM = _build_transform()
 
 
-def classify(image: Image.Image) -> dict[str, float]:
+LOW_CONFIDENCE_THRESHOLD = 0.5
+
+
+def classify(image: Image.Image) -> tuple[dict[str, float], str]:
     if image is None:
-        return {}
+        return {}, ""
     rgb = image.convert("RGB")
     x = TRANSFORM(rgb).unsqueeze(0)
     with torch.inference_mode():
         logits = MODEL(x)
         probs = torch.softmax(logits, dim=1).squeeze(0).numpy()
-    return {CLASS_NAMES[IDX_TO_CLASS[i]]: float(probs[i]) for i in range(4)}
+    label_probs = {CLASS_NAMES[IDX_TO_CLASS[i]]: float(probs[i]) for i in range(4)}
+
+    top_label, top_prob = max(label_probs.items(), key=lambda kv: kv[1])
+    notes: list[str] = []
+    if top_prob < LOW_CONFIDENCE_THRESHOLD:
+        notes.append(
+            f"⚠️ **Low confidence** (top class {top_label} at {top_prob:.0%}). "
+            "This often means the photo isn't a clear close-up of a wound, or the wound type "
+            "isn't one of the four the model was trained on. The model has no \"not a wound\" "
+            "option — it will always pick one of D / P / S / V even when the image isn't a wound."
+        )
+    if top_label == "Pressure ulcer":
+        notes.append(
+            "⚠️ **Pressure-class predictions are the model's weak spot** "
+            "(~0.41 accuracy on the held-out test set). Treat this prediction with extra skepticism."
+        )
+    return label_probs, "\n\n".join(notes)
 
 
 DESCRIPTION = """\
-Chronic wound photo classifier — predicts which of four wound types an uploaded image shows
-(diabetic ulcer, pressure ulcer, surgical wound, or venous ulcer).
+Upload a close-up photo of a chronic wound and the model returns its best guess at the wound type
+(diabetic ulcer, pressure ulcer, surgical wound, or venous ulcer) with per-class probabilities.
+
+**Before you upload — a few things worth knowing:**
+
+- **Frame the wound in the center.** The model resizes and center-crops to 224×224 pixels, so a wound
+  off in the corner of a high-resolution phone photo will get cropped out. Close-up clinical-style
+  framing (the wound fills most of the frame) gives the cleanest predictions.
+- **JPEG or PNG works fine.** Anything PIL can open.
+- **There is no "not a wound" option.** Upload a face, a sandwich, a sunset — the model will
+  confidently place it into one of the four wound classes anyway. The low-confidence warning below
+  is the only signal you'll get that the image probably isn't what the model was trained on.
+- **Pressure-ulcer predictions are weakest** (~41% accuracy on the held-out test set). When the model
+  says "Pressure ulcer," treat the call with extra skepticism.
 
 **This is a research demo, not a medical device.** It does not diagnose, triage, or replace clinician
-judgement. It was trained on the public AZH Chronic Wound Database (Anisuzzaman et al. 2022) using
+judgement. Trained on the public AZH Chronic Wound Database (Anisuzzaman et al. 2022) using
 EfficientNet-B0 with patient-grouped 10-fold cross-validation; the deployed checkpoint
 (`cv_baseline_fold5_best.pt`) reaches **0.8152 top-1** on the held-out AZH Test set (n=184).
-Pressure ulcers are the underperforming class (~0.41 accuracy) — interpret pressure-class predictions
-with extra care.
-
-Upload a wound photograph below to see the model's per-class probabilities.
 """
 
 ARTICLE = """\
@@ -110,8 +137,11 @@ location by deep neural network*, Sci. Rep. 12:20057.
 
 demo = gr.Interface(
     fn=classify,
-    inputs=gr.Image(type="pil", label="Wound photograph"),
-    outputs=gr.Label(num_top_classes=4, label="Predicted wound type"),
+    inputs=gr.Image(type="pil", label="Wound photograph (close-up, centered)"),
+    outputs=[
+        gr.Label(num_top_classes=4, label="Predicted wound type"),
+        gr.Markdown(label="Notes"),
+    ],
     title="Chronic Wound Classifier — 4-class AZH demo",
     description=DESCRIPTION,
     article=ARTICLE,