backend/severity_model.py

"""
Severity model for DermaConnect.

This module exposes a single function, `assess_severity(image_path)`, that any
part of the app can call. Today it runs a lightweight, on-device heuristic that
needs no API key, no GPU, and nothing heavier than Pillow, so the whole app is
runnable out of the box. The return shape is the contract every caller depends
on, so you can swap the guts for a real model (a trained CNN, or a hosted vision
LLM) without touching the backend.

Return contract:
{
    "score":     float,   # 0-100, higher = more severe
    "category":  str,     # one of CATEGORIES
    "confidence":float,   # 0-1, model's own confidence
    "signals":   dict,    # raw features, for debugging / clinician review
    "model":     str,     # model identifier, stored with every reading
}

--- How to swap in a real model -------------------------------------------------
1. Trained CNN (recommended for production):
   - Train / fine-tune on a labelled derm dataset (e.g. severity grades for
     eczema (EASI), psoriasis (PASI), acne (IGA)). Keep one model per condition.
   - Replace `_heuristic_assess` with a function that loads your weights once
     (module-level, not per call) and returns the same dict.

2. Hosted vision LLM (fastest to wire up, good for a demo):
   - Send the image to a multimodal model and ask for a 0-100 severity grade
     plus a one-line rationale. Map its answer into the dict below.
   - Keep the `model` field accurate so readings stay auditable.

Either way, do NOT change the return shape. The backend, the dashboards, and the
stored history all rely on it.
"""

from __future__ import annotations

import os
from typing import Dict

CATEGORIES = ["Clear", "Mild", "Moderate", "Severe"]
MODEL_ID = "heuristic-erythema-v0"


def _category_for_score(score: float) -> str:
    if score < 20:
        return "Clear"
    if score < 45:
        return "Mild"
    if score < 70:
        return "Moderate"
    return "Severe"


def _heuristic_assess(image_path: str) -> Dict:
    """
    A transparent, dependency-light stand-in for a real severity model.

    It estimates inflammation from three image signals that loosely track what a
    clinician eyeballs: redness (erythema), how much of the frame looks inflamed
    (extent), and texture irregularity (a rough proxy for scaling / lesions).
    Implemented with Pillow only, so it runs anywhere Python runs.

    This is NOT a medical device and is NOT diagnostic. It exists so the app runs
    end-to-end today. Replace it with a validated model before any real use.
    """
    from PIL import Image, ImageFilter, ImageStat

    img = Image.open(image_path).convert("RGB")
    img.thumbnail((200, 200))
    pixels = list(img.getdata())
    n = len(pixels) or 1

    # Per-pixel "red dominance". Healthy skin is ALREADY red-dominant, so the raw
    # value is useless on its own. The signal that matters is how much a pixel's
    # redness exceeds this patient's own healthy-skin baseline (local contrast).
    reds = []
    for r, g, b in pixels:
        brightness = (r + g + b) / 3.0 + 1e-6
        reds.append((r - (g + b) / 2.0) / brightness)

    # Baseline = median redness across the frame (the skin the lesion sits on).
    baseline = sorted(reds)[n // 2]
    margin = 0.06  # ignore small natural variation

    excess_sum = 0.0
    inflamed = 0
    for v in reds:
        excess = v - baseline - margin
        if excess > 0:
            excess_sum += excess
            if excess > 0.05:
                inflamed += 1

    erythema = excess_sum / n           # avg redness ABOVE healthy baseline
    extent = inflamed / n               # fraction clearly inflamed vs. own skin

    # Texture: mean edge energy via Pillow's edge filter (scaling / lesion proxy)
    edges = img.convert("L").filter(ImageFilter.FIND_EDGES)
    texture = ImageStat.Stat(edges).mean[0] / 255.0

    # Blend the signals into a 0-100 score. Weights are deliberately simple and
    # documented so a clinician can reason about them; a real model learns these.
    # Erythema dominates; extent and texture add lift. Calibrated so faint redness
    # lands Clear/Mild and dense inflammation lands Severe across the 0-100 band.
    raw = 1.7 * erythema + 0.5 * extent + 0.4 * texture
    score = float(max(0.0, min(100.0, raw * 100.0)))

    confidence = round(0.4 + 0.3 * min(1.0, extent * 3), 2)

    return {
        "score": round(score, 1),
        "category": _category_for_score(score),
        "confidence": confidence,
        "signals": {
            "erythema": round(erythema, 4),
            "extent": round(extent, 4),
            "texture": round(texture, 4),
            "baseline": round(baseline, 4),
        },
        "model": MODEL_ID,
    }


def assess_severity(image_path: str) -> Dict:
    """Public entry point. Swap the body to change models; keep the return shape."""
    if not os.path.exists(image_path):
        raise FileNotFoundError(image_path)
    return _heuristic_assess(image_path)