Add pipeline code, PWAT models, and Gradio app

README.md

@@ -1,12 +1,105 @@
----
-title: WoundNetB7 DFU Analysis
-emoji: 💻
-colorFrom: indigo
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.11.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# WoundNetB7 — Automated DFU Assessment Pipeline
+
+End-to-end pipeline for Diabetic Foot Ulcer (DFU) analysis:
+
+**Image** -> **Multiclass Segmentation** -> **PWAT Scoring** + **Fitzpatrick/ITA Estimation**
+
+## Pipeline
+
+```
+Input Image (DFU photograph)
+    |
+    v
+[1] WoundNetB7 Segmentation (EfficientNet-B7 + ASPP + CBAM + TAM)
+    -> 4-class masks: background, foot, perilesion, ulcer
+    -> Ulcer Dice: 0.927 (95% CI: [0.917, 0.936])
+    |
+    +---> [2] PWAT Estimation (XGBoost per item)
+    |         -> Items 3-8 scores (0-4 ordinal)
+    |         -> 5-fold CV: MAE 0.61-0.87, Adjacent Match 75-89%
+    |
+    +---> [3] Fitzpatrick/ITA Estimation
+    |         -> Healthy skin = foot - perilesion - ulcer
+    |         -> ITA angle + Fitzpatrick I-VI classification
+    |         -> Calibrated on 61 DFU images (86.9% exact match)
+    |
+    +---> [4] Debiasing
+              -> PWAT scores adjusted by Fitzpatrick type
+              -> 18% max group gap reduction (p < 10^-27)
+```
+
+## Quick Start
+
+```python
+from pipeline import WoundNetB7Pipeline
+
+pipe = WoundNetB7Pipeline("models/")
+result = pipe.analyze("path/to/dfu_image.png")
+print(result.summary())
+```
+
+## Models
+
+| Component | Architecture | Metric | Value |
+|-----------|-------------|--------|-------|
+| Segmentation | EfficientNet-B7 + UNet + ASPP + CBAM + TAM | Ulcer Dice | 0.927 |
+| PWAT Item 3 | XGBoost + SMOTE | MAE / Adjacent | 0.87 / 74.9% |
+| PWAT Item 7 | XGBoost + SMOTE | MAE / Adjacent | 0.61 / 89.3% |
+| Fitzpatrick | ITA calibrated thresholds | Exact match | 86.9% |
+
+## Fitzpatrick/ITA Calibrated Thresholds
+
+| Type | ITA Range | Description |
+|------|-----------|-------------|
+| I | > 46.86 | Very Light |
+| II | 34.25 - 46.86 | Light |
+| III | 20.87 - 34.25 | Intermediate |
+| IV | 3.57 - 20.87 | Tan |
+| V | -28.38 - 3.57 | Brown |
+| VI | < -28.38 | Dark |
+
+## PWAT Debiasing Factors
+
+Darker skin tones tend to overestimate tissue severity (especially Item 8: Periulcer Skin).
+Correction factors reduce this bias:
+
+| Type | P3 | P4 | P5 | P6 | P7 | P8 |
+|------|-----|-----|-----|-----|-----|-----|
+| I-II | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 |
+| III | 0.0 | 0.0 | 0.0 | 0.0 | 0.0 | -0.1 |
+| IV | -0.1 | -0.1 | 0.0 | 0.0 | 0.0 | -0.3 |
+| V | -0.2 | -0.2 | -0.1 | 0.0 | 0.0 | -0.6 |
+| VI | -0.3 | -0.3 | -0.2 | -0.1 | 0.0 | -0.9 |
+
+## Dataset
+
+- **Segmentation training**: 2,245 images (1,571 train / 336 val / 338 test)
+- **Multiclass validation**: 461 images with 5-class expert masks
+- **PWAT labels**: 1,321 images (920 expert GT + 401 AI consensus)
+- **Fitzpatrick calibration**: 61 images with expert skin type annotations
+
+## Deploy to Hugging Face
+
+```bash
+# Clone this repo to a HF Space
+git clone https://huggingface.co/spaces/YOUR_USER/woundnetb7
+# Copy models/ directory (or use git-lfs for large files)
+git lfs install
+git lfs track "*.pt" "*.pkl"
+git add . && git commit -m "Initial deployment"
+git push
+```
+
+## License
+
+Research use only. Clinical deployment requires regulatory approval.
+
+## Citation
+
+```bibtex
+@article{marquez2026woundnetb7,
+  title={WoundNetB7: Automated PWAT Protocol using Topological AI for Diabetic Foot Ulcers},
+  author={M{\'a}rquez-Sandoval, Marcelo},
+  year={2026}
+}
+```

app.py

@@ -0,0 +1,103 @@
+"""Gradio app for WoundNetB7 DFU Analysis — Hugging Face Spaces deployment.
+
+Launch locally: python app.py
+Deploy to HF: push this repo to a Hugging Face Space (GPU recommended).
+"""
+import gradio as gr
+import numpy as np
+import cv2
+import json
+from pipeline import WoundNetB7Pipeline
+
+# Initialize pipeline (loads all models once)
+pipe = WoundNetB7Pipeline(models_dir="models", use_tta=True)
+
+CLASS_COLORS_RGB = {
+    0: (0, 0, 0),
+    1: (0, 255, 0),      # Foot: green
+    2: (255, 165, 0),     # Perilesion: orange
+    3: (255, 0, 0),       # Ulcer: red
+}
+
+FITZ_COLORS = {
+    "I": "#fef3c7", "II": "#fde68a", "III": "#fbbf24",
+    "IV": "#b45309", "V": "#78350f", "VI": "#451a03",
+}
+
+
+def analyze_image(image):
+    """Main analysis function called by Gradio."""
+    if image is None:
+        return None, "Please upload an image.", "{}"
+
+    # Gradio provides RGB numpy array
+    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    # Run pipeline
+    result = pipe.analyze(img_bgr, use_tta=True)
+
+    # Create overlay visualization
+    overlay = pipe.visualize(img_bgr, result)
+
+    # Build text summary
+    summary = result.summary()
+
+    # Build JSON output
+    json_out = json.dumps(result.to_dict(), indent=2, ensure_ascii=False)
+
+    return overlay, summary, json_out
+
+
+# Build Gradio interface
+with gr.Blocks(
+    title="WoundNetB7 DFU Analysis",
+    theme=gr.themes.Soft(),
+) as demo:
+    gr.Markdown(
+        """
+        # WoundNetB7 — Diabetic Foot Ulcer Analysis
+
+        Upload a DFU image to get:
+        1. **Multiclass segmentation** (foot / perilesion / ulcer)
+        2. **Fitzpatrick skin type** via calibrated ITA (86.9% accuracy)
+        3. **PWAT scores** items 3-8 with Fitzpatrick debiasing
+
+        > Model: EfficientNet-B7 + ASPP + CBAM + TAM | Ulcer Dice: 0.927
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="DFU Image", type="numpy")
+            analyze_btn = gr.Button("Analyze", variant="primary", size="lg")
+
+        with gr.Column(scale=1):
+            output_overlay = gr.Image(label="Segmentation Overlay")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            output_text = gr.Textbox(label="Analysis Summary", lines=25, max_lines=40)
+        with gr.Column(scale=1):
+            output_json = gr.Code(label="JSON Output", language="json")
+
+    analyze_btn.click(
+        fn=analyze_image,
+        inputs=[input_image],
+        outputs=[output_overlay, output_text, output_json],
+    )
+
+    gr.Markdown(
+        """
+        ---
+        **Legend:** Green = Foot | Orange = Perilesion | Red = Ulcer
+
+        **PWAT Items:** 3=Necrotic Type, 4=Necrotic Amount, 5=Granulation Type,
+        6=Granulation Amount, 7=Edges, 8=Periulcer Skin (0=best, 4=worst)
+
+        **Debiasing:** Scores adjusted by Fitzpatrick type to reduce skin-tone bias
+        (18% max group gap reduction, p < 10^-27).
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=False)

models/pwat/xgb_pwat3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94c73bf442d9605e557d5d28578efb3b8858ffbcd8df42033685cb5a9bf5b57e
+size 1428233

models/pwat/xgb_pwat4.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7f74559882bb03542d9126386cec6c184cf50eb49340bbcfc2e197bed85cf44
+size 1431895

models/pwat/xgb_pwat5.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70ebfd69b4f41d5ef7c33b09b456cdbc194b4f5f37f8ddffcbb0050b0c8344f3
+size 1856653

models/pwat/xgb_pwat6.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69b63396d339a95eeb853daf0fc8b60f02a3c7065b8ef4f879bcab2833d6e4b4
+size 1650293

models/pwat/xgb_pwat7.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d32c766bfd592d643632425fe422826551c9d093d34b241237f986980b528cc9
+size 1556397

models/pwat/xgb_pwat8.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1dd0c58d57755d4459d0ceb71a2c2a8fd858735334516cba8e0d948ea5f2d9e
+size 1570379

pipeline.py

@@ -0,0 +1,169 @@
+"""WoundNetB7 End-to-End Pipeline: Image -> Segmentation -> PWAT + Fitzpatrick/ITA.
+
+Usage:
+    from pipeline import WoundNetB7Pipeline
+    pipe = WoundNetB7Pipeline("models/")
+    result = pipe.analyze("path/to/dfu_image.png")
+    print(result)
+"""
+import torch
+import numpy as np
+import cv2
+from pathlib import Path
+from dataclasses import dataclass, field, asdict
+from typing import Optional
+
+from src.segmentation import load_segmentation_model, segment, CLASS_NAMES, CLASS_COLORS
+from src.fitzpatrick_estimator import estimate_fitzpatrick, FitzpatrickResult
+from src.pwat_estimator import PWATPredictor, PWATResult, ITEM_NAMES
+
+
+@dataclass
+class AnalysisResult:
+    """Complete DFU analysis result from a single image."""
+    # Segmentation
+    class_distribution: dict = field(default_factory=dict)  # {class_name: percentage}
+    # Fitzpatrick
+    fitzpatrick: Optional[FitzpatrickResult] = None
+    # PWAT
+    pwat: Optional[PWATResult] = None
+    # Metadata
+    image_size: tuple = (0, 0)
+    device: str = "cpu"
+
+    def summary(self) -> str:
+        lines = ["=" * 50, "WoundNetB7 DFU Analysis", "=" * 50]
+        lines.append(f"Image: {self.image_size[1]}x{self.image_size[0]}")
+        lines.append(f"Device: {self.device}")
+
+        lines.append("\n--- Segmentation ---")
+        for cls, pct in self.class_distribution.items():
+            lines.append(f"  {cls:<15s}: {pct:5.1f}%")
+
+        if self.fitzpatrick:
+            f = self.fitzpatrick
+            lines.append("\n--- Fitzpatrick / ITA ---")
+            lines.append(f"  Type:       {f.fitzpatrick_type} ({f.fitzpatrick_label})")
+            lines.append(f"  ITA:        {f.ita_angle:.1f} +/- {f.ita_std:.1f}")
+            lines.append(f"  L* mean:    {f.l_skin_mean:.1f}")
+            lines.append(f"  Confidence: {f.confidence:.2f}")
+            lines.append(f"  Pixels:     {f.healthy_pixels:,}")
+
+        if self.pwat and self.pwat.scores_raw:
+            p = self.pwat
+            lines.append("\n--- PWAT Scores (Items 3-8) ---")
+            lines.append(f"  {'Item':<22s} {'Raw':>4s} {'Adj':>5s}")
+            lines.append("  " + "-" * 33)
+            for item in [3, 4, 5, 6, 7, 8]:
+                name = ITEM_NAMES.get(item, f"Item {item}")
+                raw = p.scores_raw.get(item, "-")
+                adj = p.scores_adjusted.get(item, "-")
+                lines.append(f"  {name:<22s} {raw:>4}  {adj:>5.1f}")
+            lines.append(f"  {'TOTAL':<22s} {p.total_raw:>4}  {p.total_adjusted:>5.1f}")
+            lines.append(f"  Fitzpatrick correction: {p.fitzpatrick_type}")
+
+        return "\n".join(lines)
+
+    def to_dict(self) -> dict:
+        d = {"image_size": self.image_size, "device": self.device, "class_distribution": self.class_distribution}
+        if self.fitzpatrick:
+            d["fitzpatrick"] = asdict(self.fitzpatrick)
+        if self.pwat:
+            d["pwat"] = {
+                "scores_raw": self.pwat.scores_raw,
+                "scores_adjusted": self.pwat.scores_adjusted,
+                "total_raw": self.pwat.total_raw,
+                "total_adjusted": self.pwat.total_adjusted,
+            }
+        return d
+
+
+class WoundNetB7Pipeline:
+    """End-to-end DFU analysis pipeline.
+
+    Args:
+        models_dir: Path to models/ directory containing:
+            - segmentation/WoundNetB7_proposed_best.pt
+            - pwat/xgb_pwat{3-8}.pkl
+        device: "cuda" or "cpu" (auto-detected if None)
+        use_tta: Use test-time augmentation for segmentation (slower but more accurate)
+    """
+
+    def __init__(self, models_dir: str = "models", device: Optional[str] = None, use_tta: bool = True):
+        self.models_dir = Path(models_dir)
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.use_tta = use_tta
+
+        # Load segmentation model
+        seg_path = self.models_dir / "segmentation" / "WoundNetB7_proposed_best.pt"
+        self.seg_model = load_segmentation_model(str(seg_path), self.device)
+        print(f"Segmentation model loaded ({sum(p.numel() for p in self.seg_model.parameters()) / 1e6:.1f}M params)")
+
+        # Load PWAT predictor
+        pwat_path = self.models_dir / "pwat"
+        self.pwat_predictor = PWATPredictor(str(pwat_path))
+        print(f"PWAT models loaded ({len(self.pwat_predictor.models)} items)")
+
+        print(f"Pipeline ready on {self.device}")
+
+    def analyze(self, image_input, use_tta: Optional[bool] = None) -> AnalysisResult:
+        """Analyze a DFU image end-to-end.
+
+        Args:
+            image_input: file path (str/Path), BGR numpy array, or RGB numpy array
+            use_tta: Override TTA setting for this call
+
+        Returns:
+            AnalysisResult with segmentation, Fitzpatrick, and PWAT data.
+        """
+        tta = use_tta if use_tta is not None else self.use_tta
+
+        # Load image
+        if isinstance(image_input, (str, Path)):
+            img_bgr = cv2.imread(str(image_input))
+            if img_bgr is None:
+                raise FileNotFoundError(f"Cannot read image: {image_input}")
+        elif isinstance(image_input, np.ndarray):
+            img_bgr = image_input if image_input.shape[2] == 3 else cv2.cvtColor(image_input, cv2.COLOR_RGB2BGR)
+        else:
+            raise TypeError(f"Unsupported input type: {type(image_input)}")
+
+        h, w = img_bgr.shape[:2]
+
+        # Step 1: Segmentation
+        seg = segment(self.seg_model, img_bgr, self.device, use_tta=tta)
+        classmap = seg["classmap"]
+
+        class_dist = {}
+        for cid, name in CLASS_NAMES.items():
+            class_dist[name] = round(float(np.mean(classmap == cid) * 100), 1)
+
+        # Step 2: Fitzpatrick estimation (from healthy skin)
+        fitz = estimate_fitzpatrick(img_bgr, seg["masks"])
+
+        # Step 3: PWAT prediction (from ulcer mask)
+        ulcer_mask = (classmap == 3).astype(np.uint8) * 255
+        pwat = self.pwat_predictor.predict(img_bgr, ulcer_mask, fitzpatrick_type=fitz.fitzpatrick_type)
+
+        return AnalysisResult(
+            class_distribution=class_dist,
+            fitzpatrick=fitz,
+            pwat=pwat,
+            image_size=(h, w),
+            device=str(self.device),
+        )
+
+    def visualize(self, img_bgr: np.ndarray, result: AnalysisResult) -> np.ndarray:
+        """Create overlay visualization of segmentation result."""
+        h, w = img_bgr.shape[:2]
+        seg = segment(self.seg_model, img_bgr, self.device, use_tta=False)
+        classmap = seg["classmap"]
+
+        overlay = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32)
+        for cid, color in CLASS_COLORS.items():
+            if cid == 0:
+                continue
+            mask = classmap == cid
+            overlay[mask] = overlay[mask] * 0.5 + np.array(color, dtype=np.float32) * 0.5
+
+        return overlay.astype(np.uint8)

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.0.0
+torchvision>=0.15.0
+segmentation-models-pytorch>=0.4.0
+timm>=1.0.3
+opencv-python-headless>=4.8.0
+numpy>=1.24.0
+pandas>=2.0.0
+scikit-learn>=1.3.0
+xgboost>=2.0.0
+joblib>=1.3.0
+gradio>=4.0.0

src/__init__.py

@@ -0,0 +1 @@
+"""WoundNetB7 DFU Analysis Pipeline — Segmentation + PWAT + Fitzpatrick/ITA."""

src/fitzpatrick_estimator.py

@@ -0,0 +1,137 @@
+"""Fitzpatrick skin type estimation via ITA (Individual Typology Angle).
+
+Calibrated on 61 DFU images with expert ground truth.
+Validation: 86.9% exact match, 98.4% adjacent, r=0.975.
+"""
+import numpy as np
+import cv2
+from dataclasses import dataclass
+from typing import Optional
+
+# Calibrated ITA thresholds for DFU clinical photography (61-image validation)
+ITA_THRESHOLDS = {
+    "I":   (46.86, float("inf")),
+    "II":  (34.25, 46.86),
+    "III": (20.87, 34.25),
+    "IV":  (3.57, 20.87),
+    "V":   (-28.38, 3.57),
+    "VI":  (float("-inf"), -28.38),
+}
+
+FITZPATRICK_LABELS = {
+    "I": "Very Light", "II": "Light", "III": "Intermediate",
+    "IV": "Tan", "V": "Brown", "VI": "Dark",
+}
+
+
+@dataclass
+class FitzpatrickResult:
+    fitzpatrick_type: str          # "I" .. "VI"
+    fitzpatrick_int: int           # 1 .. 6
+    fitzpatrick_label: str         # "Very Light" .. "Dark"
+    ita_angle: float               # ITA in degrees
+    ita_std: float                 # ITA standard deviation
+    l_skin_mean: float             # Mean L* of healthy skin
+    b_skin_mean: float             # Mean b* of healthy skin
+    healthy_pixels: int            # Number of healthy skin pixels used
+    healthy_ratio: float           # Healthy pixels / total image pixels
+    confidence: float              # 0-1 confidence score
+
+
+def compute_ita(l_values: np.ndarray, b_values: np.ndarray) -> tuple:
+    """Compute ITA angle from L* and b* values with robust trimming.
+
+    ITA = arctan((L* - 50) / b*) * (180 / pi)
+    Higher ITA = lighter skin, lower ITA = darker skin.
+    """
+    ita_per_pixel = np.degrees(np.arctan2(l_values - 50.0, b_values))
+    # Robust trimming: 5th-95th percentile
+    p5, p95 = np.percentile(ita_per_pixel, [5, 95])
+    trimmed = ita_per_pixel[(ita_per_pixel >= p5) & (ita_per_pixel <= p95)]
+    if len(trimmed) < 10:
+        trimmed = ita_per_pixel
+    return float(np.mean(trimmed)), float(np.std(trimmed))
+
+
+def classify_fitzpatrick(ita: float) -> tuple:
+    """Classify ITA angle into Fitzpatrick type using calibrated DFU thresholds."""
+    for ftype, (lo, hi) in ITA_THRESHOLDS.items():
+        if lo <= ita < hi:
+            idx = list(ITA_THRESHOLDS.keys()).index(ftype) + 1
+            return ftype, idx
+    return "III", 3  # Default fallback
+
+
+def estimate_fitzpatrick(
+    img_bgr: np.ndarray,
+    masks: dict,
+    periulcer_dilation_px: int = 40,
+) -> FitzpatrickResult:
+    """Estimate Fitzpatrick type from a DFU image using segmentation masks.
+
+    Strategy: Healthy skin = foot region - perilesion zone - ulcer.
+    ITA is computed on the healthy skin pixels only.
+
+    Args:
+        img_bgr: BGR image (H, W, 3)
+        masks: dict with keys 'foot', 'perilesion', 'ulcer' (bool arrays H, W)
+        periulcer_dilation_px: Extra dilation around wound for safety margin
+    """
+    h, w = img_bgr.shape[:2]
+    foot = masks.get("foot", np.ones((h, w), dtype=bool))
+    peri = masks.get("perilesion", np.zeros((h, w), dtype=bool))
+    ulcer = masks.get("ulcer", np.zeros((h, w), dtype=bool))
+
+    # Dilate ulcer+perilesion for safety margin
+    exclusion = (peri | ulcer).astype(np.uint8)
+    if periulcer_dilation_px > 0:
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (periulcer_dilation_px, periulcer_dilation_px))
+        exclusion = cv2.dilate(exclusion, kernel)
+    exclusion = exclusion.astype(bool)
+
+    # Healthy skin = foot minus exclusion zone
+    healthy = foot & ~exclusion
+    healthy_pixels = int(np.sum(healthy))
+
+    if healthy_pixels < 100:
+        # Fallback: use all foot pixels minus wound
+        healthy = foot & ~ulcer
+        healthy_pixels = int(np.sum(healthy))
+
+    if healthy_pixels < 50:
+        return FitzpatrickResult(
+            fitzpatrick_type="III", fitzpatrick_int=3,
+            fitzpatrick_label="Intermediate",
+            ita_angle=0.0, ita_std=0.0,
+            l_skin_mean=0.0, b_skin_mean=0.0,
+            healthy_pixels=healthy_pixels,
+            healthy_ratio=healthy_pixels / (h * w),
+            confidence=0.0,
+        )
+
+    # Convert to L*a*b*
+    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2Lab).astype(np.float32)
+    l_values = lab[healthy, 0] * (100.0 / 255.0)  # OpenCV L* is 0-255 -> 0-100
+    b_values = lab[healthy, 2] - 128.0             # OpenCV b* is 0-255 -> -128 to +127
+
+    ita_mean, ita_std = compute_ita(l_values, b_values)
+    ftype, fint = classify_fitzpatrick(ita_mean)
+
+    # Confidence: pixel count + ITA consistency + coverage
+    pixel_conf = min(healthy_pixels / 5000.0, 1.0)
+    ita_conf = max(0.0, 1.0 - (ita_std / 30.0))
+    coverage_conf = min((healthy_pixels / (h * w)) / 0.15, 1.0)
+    confidence = pixel_conf * 0.3 + ita_conf * 0.4 + coverage_conf * 0.3
+
+    return FitzpatrickResult(
+        fitzpatrick_type=ftype,
+        fitzpatrick_int=fint,
+        fitzpatrick_label=FITZPATRICK_LABELS[ftype],
+        ita_angle=round(ita_mean, 2),
+        ita_std=round(ita_std, 2),
+        l_skin_mean=round(float(np.mean(l_values)), 2),
+        b_skin_mean=round(float(np.mean(b_values)), 2),
+        healthy_pixels=healthy_pixels,
+        healthy_ratio=round(healthy_pixels / (h * w), 4),
+        confidence=round(confidence, 3),
+    )

src/pwat_estimator.py

@@ -0,0 +1,183 @@
+"""PWAT (Photographic Wound Assessment Tool) estimation — Items 3-8.
+
+Uses segmentation masks to extract color, tissue, morphological, and texture
+features, then predicts ordinal PWAT scores (0-4) via XGBoost classifiers.
+
+Includes Fitzpatrick-aware debiasing correction factors.
+"""
+import numpy as np
+import cv2
+import joblib
+from dataclasses import dataclass, field
+from typing import Optional
+from pathlib import Path
+
+ITEMS = [3, 4, 5, 6, 7, 8]
+ITEM_NAMES = {
+    3: "Necrotic Type",
+    4: "Necrotic Amount",
+    5: "Granulation Type",
+    6: "Granulation Amount",
+    7: "Edges",
+    8: "Periulcer Skin",
+}
+
+# Debiasing correction factors (calibrated from 61 DFU images)
+# Applied as: adjusted = clip(raw + factor, 0, 4)
+CORRECTION_FACTORS = {
+    "I":   {3: 0.0,  4: 0.0,  5: 0.0,  6: 0.0,  7: 0.0,  8: 0.0},
+    "II":  {3: 0.0,  4: 0.0,  5: 0.0,  6: 0.0,  7: 0.0,  8: 0.0},
+    "III": {3: 0.0,  4: 0.0,  5: 0.0,  6: 0.0,  7: 0.0,  8: -0.1},
+    "IV":  {3: -0.1, 4: -0.1, 5: 0.0,  6: 0.0,  7: 0.0,  8: -0.3},
+    "V":   {3: -0.2, 4: -0.2, 5: -0.1, 6: 0.0,  7: 0.0,  8: -0.6},
+    "VI":  {3: -0.3, 4: -0.3, 5: -0.2, 6: -0.1, 7: 0.0,  8: -0.9},
+}
+
+
+@dataclass
+class PWATResult:
+    scores_raw: dict = field(default_factory=dict)       # {item: int}
+    scores_adjusted: dict = field(default_factory=dict)  # {item: float} (debiased)
+    total_raw: int = 0
+    total_adjusted: float = 0.0
+    fitzpatrick_type: str = ""
+    features: dict = field(default_factory=dict)
+
+
+def extract_features(img_bgr: np.ndarray, ulcer_mask: np.ndarray) -> Optional[dict]:
+    """Extract 63 features from the wound region for PWAT prediction.
+
+    Features: color (RGB/HSV/Lab), tissue composition, morphology, texture.
+    """
+    b = ulcer_mask > 0 if ulcer_mask.dtype == bool else ulcer_mask > 127
+    npx = int(np.sum(b))
+    if npx < 50:
+        return None
+
+    feats = {}
+
+    # --- Color features (45) ---
+    hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV).astype(np.float32)
+    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2Lab).astype(np.float32)
+    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.float32)
+
+    for cs, arr, names in [("rgb", rgb, ["R", "G", "B"]), ("hsv", hsv, ["H", "S", "V"]), ("lab", lab, ["L", "a", "b"])]:
+        for ci, cn in enumerate(names):
+            vals = arr[b, ci]
+            feats[f"{cs}_{cn}_mean"] = float(np.mean(vals))
+            feats[f"{cs}_{cn}_std"] = float(np.std(vals))
+            feats[f"{cs}_{cn}_median"] = float(np.median(vals))
+            feats[f"{cs}_{cn}_p25"] = float(np.percentile(vals, 25))
+            feats[f"{cs}_{cn}_p75"] = float(np.percentile(vals, 75))
+
+    # --- Tissue composition (5) ---
+    h, s, v = hsv[b, 0], hsv[b, 1], hsv[b, 2]
+    l_ch = lab[b, 0] * (100 / 255)
+    a_ch = lab[b, 1] - 128
+
+    eschar = ((v < 100) & (s < 60)) | (v < 60)
+    slough = (h >= 15) & (h <= 50) & (s > 25) & (v > 70) & ~eschar
+    gran = (((h < 15) | (h > 155)) & (s > 35) & (v > 60) & (a_ch > 5)) & ~eschar
+    necro = (s < 45) & (v >= 60) & (v < 160) & (l_ch < 55) & ~eschar & ~gran
+
+    feats["tissue_gran_pct"] = float(np.sum(gran) / npx * 100)
+    feats["tissue_eschar_pct"] = float(np.sum(eschar) / npx * 100)
+    feats["tissue_slough_pct"] = float(np.sum(slough) / npx * 100)
+    feats["tissue_necro_pct"] = float(np.sum(necro) / npx * 100)
+    feats["tissue_necro_total"] = feats["tissue_eschar_pct"] + feats["tissue_slough_pct"] + feats["tissue_necro_pct"]
+
+    # --- Morphological features (7) ---
+    mask_u8 = b.astype(np.uint8) if b.dtype == bool else (ulcer_mask > 127).astype(np.uint8)
+    cnts, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if cnts:
+        cnt = max(cnts, key=cv2.contourArea)
+        area = cv2.contourArea(cnt)
+        perim = cv2.arcLength(cnt, True)
+        circ = 4 * np.pi * area / (perim**2) if perim > 0 else 0
+        feats["morph_area"] = float(area)
+        feats["morph_perimeter"] = float(perim)
+        feats["morph_circularity"] = float(circ)
+        feats["morph_irregularity"] = float(1 - circ)
+        x, y, w2, h2 = cv2.boundingRect(cnt)
+        feats["morph_aspect_ratio"] = float(w2 / (h2 + 1e-8))
+        feats["morph_extent"] = float(area / (w2 * h2 + 1e-8))
+        hull = cv2.convexHull(cnt)
+        feats["morph_solidity"] = float(area / (cv2.contourArea(hull) + 1e-8))
+
+    # --- Texture features (4) ---
+    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
+    wound_gray = gray[b]
+    feats["texture_mean"] = float(np.mean(wound_gray))
+    feats["texture_std"] = float(np.std(wound_gray))
+    hist_vals = np.histogram(wound_gray, bins=64, density=True)[0]
+    feats["texture_entropy"] = float(-np.sum(hist_vals * np.log2(hist_vals + 1e-10)))
+
+    dilated = cv2.dilate(mask_u8 * 255, np.ones((5, 5), np.uint8))
+    eroded = cv2.erode(mask_u8 * 255, np.ones((5, 5), np.uint8))
+    edge_zone = (dilated - eroded) > 127
+    if np.any(edge_zone):
+        feats["edge_gradient"] = float(np.mean(np.abs(cv2.Sobel(gray.astype(np.float32), cv2.CV_32F, 1, 0)[edge_zone])))
+
+    # --- ROI features (2) ---
+    feats["wound_npx"] = float(npx)
+    feats["wound_ratio"] = float(npx / (img_bgr.shape[0] * img_bgr.shape[1]))
+
+    return feats
+
+
+class PWATPredictor:
+    """Predicts PWAT items 3-8 from wound features using trained XGBoost models."""
+
+    def __init__(self, models_dir: str):
+        self.models = {}
+        models_path = Path(models_dir)
+        for item in ITEMS:
+            pkl = models_path / f"xgb_pwat{item}.pkl"
+            if pkl.exists():
+                self.models[item] = joblib.load(pkl)
+
+    def predict(
+        self,
+        img_bgr: np.ndarray,
+        ulcer_mask: np.ndarray,
+        fitzpatrick_type: str = "III",
+        feature_cols: Optional[list] = None,
+    ) -> PWATResult:
+        """Predict PWAT scores for a single image.
+
+        Args:
+            img_bgr: BGR image
+            ulcer_mask: Binary ulcer mask (H, W)
+            fitzpatrick_type: Fitzpatrick type for debiasing ("I" .. "VI")
+            feature_cols: Ordered feature column names (must match training order).
+                          If None, uses all extracted features sorted alphabetically.
+        """
+        feats = extract_features(img_bgr, ulcer_mask)
+        if feats is None:
+            return PWATResult(fitzpatrick_type=fitzpatrick_type)
+
+        # Build feature vector
+        if feature_cols is None:
+            feature_cols = sorted(feats.keys())
+        X = np.array([[feats.get(c, 0.0) for c in feature_cols]])
+
+        scores_raw = {}
+        scores_adj = {}
+        for item in ITEMS:
+            if item in self.models:
+                pred = int(self.models[item].predict(X)[0])
+                scores_raw[item] = pred
+                factor = CORRECTION_FACTORS.get(fitzpatrick_type, {}).get(item, 0.0)
+                scores_adj[item] = float(np.clip(pred + factor, 0, 4))
+            else:
+                scores_raw[item] = 0
+                scores_adj[item] = 0.0
+
+        return PWATResult(
+            scores_raw=scores_raw,
+            scores_adjusted=scores_adj,
+            total_raw=sum(scores_raw.values()),
+            total_adjusted=round(sum(scores_adj.values()), 1),
+            fitzpatrick_type=fitzpatrick_type,
+            features=feats,
+        )

src/segmentation.py

@@ -0,0 +1,256 @@
+"""WoundNetB7 multiclass segmentation model — 4 classes (bg, foot, perilesion, ulcer).
+
+Architecture: EfficientNet-B7 encoder + ASPP + CBAM + TAM + UNet decoder.
+Checkpoint: Track B multiclass, ulcer Dice = 0.927 (Bootstrap CI: [0.917, 0.936]).
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import segmentation_models_pytorch as smp
+import numpy as np
+import cv2
+from pathlib import Path
+
+IMG_SIZE = 512
+MEAN = np.array([0.485, 0.456, 0.406])
+STD = np.array([0.229, 0.224, 0.225])
+CLASS_NAMES = {0: "background", 1: "foot", 2: "perilesion", 3: "ulcer"}
+CLASS_COLORS = {
+    0: (0, 0, 0),
+    1: (0, 255, 0),
+    2: (255, 165, 0),
+    3: (255, 0, 0),
+}
+
+
+# ---------------------------------------------------------------------------
+# Architecture modules (match checkpoint weights exactly)
+# ---------------------------------------------------------------------------
+
+class ChannelAttention(nn.Module):
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channels // reduction, channels, bias=False),
+        )
+
+    def forward(self, x):
+        avg_out = self.mlp(x.mean(dim=[2, 3]))
+        max_out = self.mlp(x.amax(dim=[2, 3]))
+        attn = torch.sigmoid(avg_out + max_out).unsqueeze(-1).unsqueeze(-1)
+        return x * attn
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super().__init__()
+        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
+
+    def forward(self, x):
+        avg_out = x.mean(dim=1, keepdim=True)
+        max_out = x.amax(dim=1, keepdim=True)
+        attn = torch.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))
+        return x * attn
+
+
+class CBAM(nn.Module):
+    def __init__(self, channels, reduction=16, kernel_size=7):
+        super().__init__()
+        self.ca = ChannelAttention(channels, reduction)
+        self.sa = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        return self.sa(self.ca(x))
+
+
+class DifferentiableFractalDimension(nn.Module):
+    def __init__(self, scales=None):
+        super().__init__()
+        self.scales = scales or [2, 4, 8, 16, 32]
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        counts = []
+        for s in self.scales:
+            if s >= H or s >= W:
+                continue
+            pooled = F.avg_pool2d(x, kernel_size=s, stride=s)
+            n_boxes = torch.sigmoid(10.0 * (pooled - 0.1)).sum(dim=[2, 3])
+            counts.append(n_boxes)
+        if len(counts) < 2:
+            return torch.ones(B, C, device=x.device)
+        log_s = torch.log(torch.tensor([float(s) for s in self.scales[: len(counts)]], device=x.device))
+        log_c = torch.stack([torch.log(c + 1) for c in counts], dim=-1)
+        n = log_s.shape[0]
+        sx, sxx = log_s.sum(), (log_s**2).sum()
+        sy = log_c.sum(dim=-1)
+        sxy = (log_c * log_s.unsqueeze(0).unsqueeze(0)).sum(dim=-1)
+        slope = (n * sxy - sx * sy) / (n * sxx - sx**2 + 1e-8)
+        return -slope.mean(dim=1, keepdim=True).unsqueeze(-1).unsqueeze(-1)
+
+
+class DifferentiableEulerCharacteristic(nn.Module):
+    def forward(self, x):
+        B, C, H, W = x.shape
+        b = torch.sigmoid(10.0 * (torch.sigmoid(x) - 0.5))
+        V = b.sum(dim=[2, 3])
+        E_h = (b[:, :, :, :-1] * b[:, :, :, 1:]).sum(dim=[2, 3])
+        E_v = (b[:, :, :-1, :] * b[:, :, 1:, :]).sum(dim=[2, 3])
+        F_val = (b[:, :, :-1, :-1] * b[:, :, :-1, 1:] * b[:, :, 1:, :-1] * b[:, :, 1:, 1:]).sum(dim=[2, 3])
+        euler = V - E_h - E_v + F_val
+        return euler.mean(dim=1, keepdim=True).unsqueeze(-1).unsqueeze(-1) / (H * W)
+
+
+class TopologicalAttentionModule(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.fractal = DifferentiableFractalDimension()
+        self.euler = DifferentiableEulerCharacteristic()
+        self.alpha = nn.Parameter(torch.tensor(1.0))
+        self.beta = nn.Parameter(torch.tensor(1.0))
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels + 2, in_channels, 1),
+            nn.BatchNorm2d(in_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_channels, in_channels, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        fm = self.fractal(x).expand(B, 1, H, W)
+        em = self.euler(x).expand(B, 1, H, W)
+        attn = self.conv(torch.cat([x, self.alpha * fm, self.beta * em], dim=1))
+        return x * attn + x
+
+
+class ASPP(nn.Module):
+    def __init__(self, in_ch, out_ch, rates=None):
+        super().__init__()
+        rates = rates or [6, 12, 18]
+        self.conv1x1 = nn.Sequential(nn.Conv2d(in_ch, out_ch, 1), nn.BatchNorm2d(out_ch), nn.ReLU(True))
+        self.atrous = nn.ModuleList(
+            [nn.Sequential(nn.Conv2d(in_ch, out_ch, 3, padding=r, dilation=r), nn.BatchNorm2d(out_ch), nn.ReLU(True)) for r in rates]
+        )
+        self.pool = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_ch, out_ch, 1), nn.ReLU(True))
+        self.project = nn.Sequential(
+            nn.Conv2d(out_ch * (2 + len(rates)), out_ch, 1), nn.BatchNorm2d(out_ch), nn.ReLU(True), nn.Dropout(0.5)
+        )
+
+    def forward(self, x):
+        size = x.shape[2:]
+        feats = [self.conv1x1(x)] + [a(x) for a in self.atrous]
+        feats.append(F.interpolate(self.pool(x), size=size, mode="bilinear", align_corners=False))
+        return self.project(torch.cat(feats, dim=1))
+
+
+class WoundNetB7(nn.Module):
+    """WoundNetB7 matching the Track B checkpoint structure."""
+
+    NUM_CLASSES = 4
+
+    def __init__(self, num_classes=4):
+        super().__init__()
+        self.backbone = smp.Unet(encoder_name="efficientnet-b7", encoder_weights=None, in_channels=3, classes=num_classes)
+        enc_ch = self.backbone.encoder.out_channels[-1]
+        self.aspp = ASPP(enc_ch, enc_ch)
+        self.cbam = CBAM(num_classes, reduction=max(1, num_classes // 2))
+        self.tam = TopologicalAttentionModule(num_classes)
+        self.diffusion_weight = nn.Parameter(torch.tensor(0.01))
+
+    def forward(self, x):
+        features = list(self.backbone.encoder(x))
+        features[-1] = self.aspp(features[-1])
+        try:
+            dec = self.backbone.decoder(features)
+        except TypeError:
+            dec = self.backbone.decoder(*features)
+        seg = self.backbone.segmentation_head(dec)
+        seg = self.cbam(seg)
+        seg = self.tam(seg)
+        return seg
+
+
+# ---------------------------------------------------------------------------
+# Inference helpers
+# ---------------------------------------------------------------------------
+
+def preprocess(img_bgr: np.ndarray) -> torch.Tensor:
+    """BGR image -> normalized CHW tensor (1, 3, 512, 512)."""
+    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_LINEAR)
+    img = (img.astype(np.float32) / 255.0 - MEAN) / STD
+    return torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0).float()
+
+
+def tta_inference(model: nn.Module, img_tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
+    """6-fold TTA -> averaged softmax probabilities (1, C, H, W)."""
+    transforms = [
+        lambda x: x,
+        lambda x: torch.flip(x, [3]),
+        lambda x: torch.flip(x, [2]),
+        lambda x: torch.rot90(x, 1, [2, 3]),
+        lambda x: torch.rot90(x, 2, [2, 3]),
+        lambda x: torch.rot90(x, 3, [2, 3]),
+    ]
+    inverse = [
+        lambda x: x,
+        lambda x: torch.flip(x, [3]),
+        lambda x: torch.flip(x, [2]),
+        lambda x: torch.rot90(x, 3, [2, 3]),
+        lambda x: torch.rot90(x, 2, [2, 3]),
+        lambda x: torch.rot90(x, 1, [2, 3]),
+    ]
+    probs_sum = None
+    with torch.no_grad():
+        for tfm, inv in zip(transforms, inverse):
+            out = model(tfm(img_tensor).to(device))
+            if isinstance(out, (tuple, list)):
+                out = out[0]
+            if isinstance(out, dict):
+                out = out["seg"]
+            p = inv(F.softmax(out, dim=1))
+            probs_sum = p if probs_sum is None else probs_sum + p
+    return probs_sum / len(transforms)
+
+
+def load_segmentation_model(checkpoint_path: str, device: torch.device) -> nn.Module:
+    """Load WoundNetB7 from checkpoint."""
+    model = WoundNetB7(num_classes=4)
+    state = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    # Remove PWAT head keys if present
+    state = {k: v for k, v in state.items() if not k.startswith("pwat_head.")}
+    model.load_state_dict(state, strict=False)
+    model.to(device).eval()
+    return model
+
+
+def segment(model: nn.Module, img_bgr: np.ndarray, device: torch.device, use_tta: bool = True) -> dict:
+    """Run segmentation on a BGR image.
+
+    Returns dict with:
+        classmap: (H, W) uint8 with class indices 0-3
+        masks: dict of per-class binary masks {cls_name: (H, W) bool}
+        probs: (4, H, W) float32 softmax probabilities
+    """
+    h, w = img_bgr.shape[:2]
+    tensor = preprocess(img_bgr)
+
+    if use_tta:
+        probs = tta_inference(model, tensor, device)
+    else:
+        with torch.no_grad():
+            out = model(tensor.to(device))
+            if isinstance(out, (tuple, list)):
+                out = out[0]
+            if isinstance(out, dict):
+                out = out["seg"]
+            probs = F.softmax(out, dim=1)
+
+    probs_np = probs[0].cpu().numpy()
+    probs_resized = np.stack([cv2.resize(probs_np[c], (w, h), interpolation=cv2.INTER_LINEAR) for c in range(4)])
+    classmap = probs_resized.argmax(axis=0).astype(np.uint8)
+    masks = {name: (classmap == cid) for cid, name in CLASS_NAMES.items() if cid > 0}
+    return {"classmap": classmap, "masks": masks, "probs": probs_resized}