app.py

"""
Skin Disease Detection — Demo
------------------------------
A Gradio app for Hugging Face Spaces with two tabs:
  1. Image Diagnosis — upload (or pick a test image) and get a classifier
     prediction plus a plain-language, urgency-aware explanation.
  2. Chat Assistant — a multi-turn triage chatbot. It can accept an image,
     asks the kind of follow-up questions a real triage intake would, and
     after enough context gives a preliminary assessment with urgency-tiered
     guidance.

IMPORTANT: This remains a research/portfolio project, not a validated
medical device. It has not been clinically evaluated or regulatory
cleared. It should never be the sole basis for a health decision — that
framing shows up once, clearly, rather than as a repeated warning block
after every message, but it's still true and still matters.
"""

import os

import gradio as gr
import spaces
import requests
from PIL import Image
from transformers import pipeline

# ---------------------------------------------------------------------------
# Config — override any of these via Space "Variables and secrets"
# ---------------------------------------------------------------------------
IMAGE_MODEL_ID = os.environ.get("IMAGE_MODEL_ID", "Anwarkh1/Skin_Cancer-Image_Classification")
# Hugging Face retired the old api-inference.huggingface.co Serverless API in
# favor of the "Inference Providers" router (OpenAI-compatible). Any chat
# model listed at https://huggingface.co/models?inference_provider=hf-inference
# works here.
CHAT_MODEL_ID = os.environ.get("CHAT_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
HF_TOKEN = os.environ.get("HF_TOKEN")  # add as a Space SECRET to enable the chatbot
ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"



# Optional second classifier specifically for pigmented-lesion cases (HAM10000-based,
# where melanoma is its own clean label). Useful because broader taxonomies like
# DermNet lump "Melanoma Skin Cancer, Nevi and Moles" into one label — this lets a
# mole/nevus-flagged image get a second, more specific opinion. Leave unset to disable.
LESION_MODEL_ID = os.environ.get("LESION_MODEL_ID", "")

# ---------------------------------------------------------------------------
# Urgency tiering — keyword-based rather than an exact-match label dict, so it
# works across different classifier taxonomies (the original 7-class HAM10000
# labels, the 23-class DermNet labels, or any future swap) without needing to
# hardcode every exact label string. Checked in priority order.
# ---------------------------------------------------------------------------
URGENT_KEYWORDS = [
    "melanoma", "malignant", "carcinoma", "actinic keratos",
    "bullous", "lupus", "systemic", "vasculitis", "cellulitis",
]
PROMPT_DOCTOR_KEYWORDS = [
    "fungal", "fungus", "candidiasis", "tinea", "ringworm",
    "scabies", "lyme", "infestation", "wart", "molluscum",
    "viral", "herpes", "hpv", "std", "exanthem", "drug eruption",
    "impetigo", "bacterial",
]
# Anything not matching the above falls through to "general_care" — typically
# manageable/chronic conditions (acne, eczema, psoriasis, hives, hair loss,
# benign tumors, common moles, etc.) where general skin-care guidance fits.

TIER_LABELS = {
    "urgent": "higher-risk — recommend seeing a dermatologist soon",
    "prompt_doctor_visit": "needs a proper diagnosis/prescription — recommend seeing a doctor, not urgent-emergency",
    "general_care": "typically manageable — general skin-care guidance fits, doctor visit optional",
}


def classify_tier(label):
    """Map a classifier label to an urgency tier via keyword matching."""
    l = label.lower()
    for kw in URGENT_KEYWORDS:
        if kw in l:
            return "urgent"
    for kw in PROMPT_DOCTOR_KEYWORDS:
        if kw in l:
            return "prompt_doctor_visit"
    return "general_care"


# Friendly one-line descriptions for labels we recognize exactly (both the
# original 7-class HAM10000 set and common DermNet-style labels). Purely
# cosmetic — if a label isn't here, the tier-based guidance still works fine
# without a description.
KNOWN_DESCRIPTIONS = {
    "actinic keratoses": "A rough, scaly patch caused by sun damage.",
    "basal cell carcinoma": "The most common type of skin cancer.",
    "melanoma": "The most serious common type of skin cancer.",
    "benign keratosis-like lesions": "Non-cancerous growths such as seborrheic keratoses or solar lentigines.",
    "dermatofibroma": "A common benign skin nodule, often on the legs.",
    "melanocytic nevi": "Ordinary moles.",
    "vascular lesions": "Blood-vessel-related marks such as angiomas.",
}

SYSTEM_PROMPT = """You are a warm, knowledgeable skin-health triage assistant. You are not a doctor and cannot diagnose anyone — but your job is to be genuinely useful, not to hide behind constant disclaimers.

CONVERSATION STYLE
- Have a real back-and-forth. Ask about one or two things at a time, not a long checklist at once.
- Useful things to learn over the conversation (don't force all of them if the user has already told you, or if it's clearly unnecessary):
  - How long they've had it / when they first noticed it
  - Whether it's changed recently in size, shape, or color
  - Symptoms: itching, bleeding, pain, crusting, oozing, discharge, fever
  - Personal or family history of relevant conditions
  - Relevant exposure history (sun, new products, contacts, travel, bites)
  - Where on the body it is
- If the user shared an image, you'll also receive the image classifier's findings as extra context, including an urgency tier (not shown verbatim to the user) — weave that in naturally rather than reading out raw percentages.

GIVING A PRELIMINARY ASSESSMENT
Once you have enough context (often after 2-5 exchanges, sooner if the user just wants a quick read, or immediately if red flags are already obvious), give a clear preliminary assessment using three tiers:
- **URGENT** (cancer-related findings, or things like bullous/autoimmune/systemic conditions, cellulitis, vasculitis) OR any red-flag symptoms (rapid growth, irregular/changing borders, multiple colors, asymmetry, a sore that won't heal, bleeding, spreading redness, fever) — be direct: recommend seeing a dermatologist or doctor soon, and explain briefly why those signs matter. Don't soften this into vague "keep an eye on it" language. Note: some categories (like a classifier label that groups "melanoma" together with ordinary moles) genuinely can't be told apart by the model — when that ambiguity exists, say so plainly and default to recommending it get checked rather than assuming it's benign.
- **NEEDS A DOCTOR VISIT, NOT URGENT** (infections — fungal, viral, bacterial, STD-related, infestations like scabies) — explain that these generally need a proper diagnosis and often a prescription to clear up, so a doctor visit is the right next step, but it's not an emergency. For anything STD-related, stay factual and non-judgmental, and emphasize in-person testing rather than guessing from a photo.
- **TYPICALLY MANAGEABLE** (acne, eczema, psoriasis, hives, hair loss, common moles, benign growths, contact dermatitis, etc.) with no red flags — explain what that category usually is, and give general skin-care guidance: gentle skincare habits, sun protection, not picking or scratching, and watching for changes. It's fine to mention a dermatologist visit is a reasonable option too, especially if they're unsure, worried, or it's not improving — but don't invent a treatment plan or prescribe anything.
- Never give specific medications, dosages, or treatment prescriptions — general skin-care habits are fine; treating a self- or AI-identified condition with anything beyond general care is not.
- You don't need to repeat a formal warning block every message. Say once, naturally, as part of your assessment — not as a bolted-on disclaimer — that this is a preliminary read from a conversation and an image classifier, and that an in-person exam is what actually confirms things. Then move on; don't repeat it every turn.

URGENT SITUATIONS
If the user describes something urgent — rapidly growing lesion, uncontrolled bleeding, signs of spreading infection, fever with a skin issue, severe pain — tell them clearly to seek in-person or emergency care promptly, regardless of how many turns you've had.

Keep replies conversational length — a few sentences for a question, a bit longer for the assessment itself."""


# ---------------------------------------------------------------------------
# Lazy-loaded models (so the Space boots fast and only loads on first use)
# ---------------------------------------------------------------------------
_classifier = None
_lesion_classifier = None


def get_classifier():
    global _classifier
    if _classifier is None:
        import torch

        device = 0 if torch.cuda.is_available() else -1
        _classifier = pipeline("image-classification", model=IMAGE_MODEL_ID, device=device)
    return _classifier


def get_lesion_classifier():
    """Optional second opinion model for pigmented-lesion cases. Returns None if disabled."""
    global _lesion_classifier
    if not LESION_MODEL_ID:
        return None
    if _lesion_classifier is None:
        import torch

        device = 0 if torch.cuda.is_available() else -1
        _lesion_classifier = pipeline("image-classification", model=LESION_MODEL_ID, device=device)
    return _lesion_classifier


def chat_completion(messages, max_tokens=500):
    """Call the Hugging Face Inference Providers router directly (OpenAI-compatible)."""
    if not HF_TOKEN:
        return None
    resp = requests.post(
        ROUTER_URL,
        headers={"Authorization": f"Bearer {HF_TOKEN}"},
        json={"model": CHAT_MODEL_ID, "messages": messages, "max_tokens": max_tokens},
        timeout=60,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["choices"][0]["message"]["content"]


def classify_image(image):
    """Run the primary classifier and return (preds, formatted_lines). If the top
    prediction looks like a mixed/ambiguous pigmented-lesion label and a second
    lesion-specific model is configured, also run that for a second opinion."""
    clf = get_classifier()
    preds = clf(image, top_k=5)
    lines = []
    for p in preds:
        label = p["label"]
        score = p["score"] * 100
        desc = KNOWN_DESCRIPTIONS.get(label.lower(), "")
        desc_txt = f"  \n  _{desc}_" if desc else ""
        lines.append(f"- **{label}** — {score:.1f}%{desc_txt}")

    top_label = preds[0]["label"].lower()
    if ("nevi" in top_label or "mole" in top_label) and "melanoma" in top_label:
        lesion_clf = get_lesion_classifier()
        if lesion_clf is not None:
            lesion_preds = lesion_clf(image, top_k=3)
            lines.append("\n_Second opinion from a lesion-specific model (melanoma is a distinct label here):_")
            for p in lesion_preds:
                lines.append(f"- **{p['label']}** — {p['score']*100:.1f}%")
            preds = preds + [{"label": f"[lesion-model] {p['label']}", "score": p["score"]} for p in lesion_preds]

    return preds, lines


def build_classifier_context(preds):
    """Turn classifier output into a compact context block for the LLM (not shown raw to the user)."""
    parts = []
    for p in preds[:4]:
        label = p["label"]
        if label.startswith("[lesion-model]"):
            # second-opinion predictions: tier by the underlying label, minus the tag
            bare = label.replace("[lesion-model] ", "")
            tier = classify_tier(bare)
            parts.append(f"{label} ({p['score']*100:.1f}%, second-opinion, {tier})")
            continue
        tier = classify_tier(label)
        desc = KNOWN_DESCRIPTIONS.get(label.lower(), "")
        parts.append(f"{label} ({p['score']*100:.1f}%, {tier}{': ' + desc if desc else ''})")
    return "Image classifier findings for your use (do not read these percentages verbatim): " + "; ".join(parts)


# ---------------------------------------------------------------------------
# Image diagnosis tab (single-shot: upload -> classify -> explain)
# ---------------------------------------------------------------------------
@spaces.GPU
def diagnose_image(image):
    if image is None:
        return "Please upload an image or pick one of the test images first."

    try:
        preds, lines_list = classify_image(image)
    except Exception as e:
        return f"**Model error:** {e}\n\nMake sure `IMAGE_MODEL_ID` is a valid image-classification model."

    lines = ["### 🔬 Classifier prediction\n"] + lines_list

    if HF_TOKEN:
        context = build_classifier_context(preds)
        prompt = (
            f"{context}\n\n"
            "The user just uploaded a single image with no conversation yet. Give a short "
            "preliminary assessment following your instructions: plain-language explanation of "
            "the top finding, urgency-tiered guidance, and a brief natural mention that an "
            "in-person dermatologist exam is what actually confirms things. Keep it to 4-6 sentences."
        )
        try:
            explanation = chat_completion(
                [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=300,
            )
            lines.append(f"\n### 🤖 Assessment\n{explanation}")
        except Exception as e:
            lines.append(f"\n_(Assessment unavailable: {e})_")
    else:
        lines.append(
            "\n_Add an `HF_TOKEN` secret to this Space to also get a plain-language "
            "assessment from the chat model here._"
        )

    lines.append(
        "\n---\n_Tip: use the **Chat Assistant** tab for a fuller triage conversation — "
        "it can ask follow-up questions and give a more tailored read._"
    )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Chat tab — multimodal, multi-turn triage
# ---------------------------------------------------------------------------
def _extract_text_and_image(message):
    """Normalize gr.ChatInterface(multimodal=True) message input into (text, image_path_or_None)."""
    if isinstance(message, dict):
        text = message.get("text", "") or ""
        files = message.get("files") or []
        image_path = files[0] if files else None
        return text, image_path
    return str(message), None


def chat_respond(message, history):
    if not HF_TOKEN:
        return (
            "Chat isn't configured yet — add an `HF_TOKEN` secret to this Space "
            "(Settings → Variables and secrets) to enable the chatbot."
        )

    text, image_path = _extract_text_and_image(message)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        if isinstance(turn, dict):
            role = turn.get("role")
            content = turn.get("content")
            if isinstance(content, str):
                messages.append({"role": role, "content": content})
        else:
            user_msg, bot_msg = turn
            if isinstance(user_msg, str):
                messages.append({"role": "user", "content": user_msg})
            if bot_msg:
                messages.append({"role": "assistant", "content": bot_msg})

    user_content = text
    if image_path:
        try:
            img = Image.open(image_path)
            preds, _ = classify_image(img)
            context = build_classifier_context(preds)
            user_content = f"{context}\n\nUser's message: {text or '(no message, just shared an image)'}"
        except Exception as e:
            user_content = f"(Image analysis failed: {e})\n\nUser's message: {text}"

    messages.append({"role": "user", "content": user_content})

    try:
        return chat_completion(messages, max_tokens=500)
    except Exception as e:
        return f"Sorry, I hit an error talking to the model: {e}"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
EXAMPLES_DIR = "examples"
example_images = []
if os.path.isdir(EXAMPLES_DIR):
    example_images = [
        os.path.join(EXAMPLES_DIR, f)
        for f in sorted(os.listdir(EXAMPLES_DIR))
        if f.lower().endswith((".jpg", ".jpeg", ".png"))
    ]

with gr.Blocks(title="Skin Disease Detection — Demo") as demo:
    gr.Markdown("# 🩺 Skin Disease Detection")

    with gr.Tab("📷 Quick Image Check"):
        with gr.Row():
            with gr.Column():
                img_in = gr.Image(type="pil", label="Upload a skin image")
                if example_images:
                    gr.Examples(examples=example_images, inputs=img_in, label="Or try a test image")
                else:
                    gr.Markdown(
                        "_No bundled test images found. Run `scripts/download_examples.py` "
                        "before deploying, or just upload your own image._"
                    )
                analyze_btn = gr.Button("Analyze image", variant="primary")
            with gr.Column():
                result_md = gr.Markdown()
        analyze_btn.click(diagnose_image, inputs=img_in, outputs=result_md)

    with gr.Tab("💬 Chat Assistant"):
        gr.ChatInterface(
            fn=chat_respond,
            multimodal=True,
            description=(
                "Talk through what you're noticing, and optionally attach a photo (📎). "
                "The assistant will ask a few follow-up questions before giving a preliminary read."
            ),
        )

    gr.Markdown(
        "---\nBuilt with 🤗 Transformers + Gradio. "
        f"Image model: `{IMAGE_MODEL_ID}`  ·  Chat model: `{CHAT_MODEL_ID}`"
    )

if __name__ == "__main__":
    demo.launch()