app.py

import functools
from dataclasses import dataclass

import gradio as gr
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor


MODEL_ID = "openai/clip-vit-base-patch32"

torch.set_num_threads(2)


@dataclass(frozen=True)
class VisualCategory:
    key: str
    label: str
    prompt: str
    why: str
    home_care: str
    watch_for: str


CATEGORIES = [
    VisualCategory(
        "rash",
        "Rash or skin irritation",
        "a close-up medical photo of a rash, skin irritation, redness, or dermatitis",
        "The image may show redness, texture change, or irritated skin.",
        "Avoid scratching, keep the area clean and dry, and consider whether a new soap, food, plant, fabric, or medication could be involved.",
        "Get medical advice if it spreads quickly, becomes very painful, involves fever, or affects the face, eyes, or breathing.",
    ),
    VisualCategory(
        "hives",
        "Hives or allergic reaction",
        "a close-up medical photo of hives, raised welts, or an allergic skin reaction",
        "The image may resemble raised or blotchy areas that can happen with allergic reactions.",
        "Remove likely triggers if known and monitor whether the reaction is changing.",
        "Seek urgent help for swelling of the lips, tongue, throat, wheezing, dizziness, or trouble breathing.",
    ),
    VisualCategory(
        "infection",
        "Possible infection signs",
        "a close-up medical photo of infected skin, increasing redness, swelling, warmth, pus, or cellulitis",
        "The image may show redness, swelling, or drainage patterns that can overlap with infection.",
        "Keep the area clean, avoid squeezing or picking, and mark the edge of redness if you need to track spread.",
        "Contact a clinician promptly if redness expands, pus appears, red streaks develop, fever occurs, or pain gets worse.",
    ),
    VisualCategory(
        "bruise",
        "Bruise or contusion",
        "a close-up medical photo of a bruise, contusion, purple discoloration, or black-and-blue mark",
        "The image may show discoloration after impact or pressure.",
        "Rest the area and use a wrapped cold pack for short periods during the first day.",
        "Get checked if the bruise is unexplained, very large, near the eye, follows major trauma, or comes with severe pain or limited movement.",
    ),
    VisualCategory(
        "sprain",
        "Swelling or possible sprain",
        "a close-up medical photo of swelling around an ankle, wrist, joint, or soft tissue injury",
        "The image may show swelling or shape change around a joint or soft tissue area.",
        "Rest, elevate, and use a wrapped cold pack. Avoid putting full weight on a painful joint until you know more.",
        "Seek care if there is deformity, numbness, inability to bear weight, severe pain, or concern for fracture.",
    ),
    VisualCategory(
        "cut",
        "Cut, scrape, or abrasion",
        "a close-up medical photo of a cut, scrape, abrasion, or broken skin",
        "The image may show broken skin or a surface wound.",
        "Rinse gently with clean water, apply light pressure for bleeding, and cover with a clean bandage.",
        "Seek care for deep wounds, animal or human bites, dirt that will not rinse out, spreading redness, or bleeding that will not stop.",
    ),
    VisualCategory(
        "burn",
        "Minor burn or heat injury",
        "a close-up medical photo of a minor burn, scald, sunburn, or heat injury on skin",
        "The image may show redness, blistering, or skin changes consistent with a burn.",
        "Cool the area under running cool water, avoid ice, and cover loosely with a clean nonstick dressing.",
        "Seek urgent care for large burns, burns on the face/hands/genitals, chemical/electrical burns, or white/charred skin.",
    ),
    VisualCategory(
        "blister",
        "Blister or friction injury",
        "a close-up medical photo of a blister, fluid-filled bump, or friction injury",
        "The image may show a fluid-filled raised area, often from rubbing, heat, or irritation.",
        "Protect the area from friction and keep the skin covering intact if possible.",
        "Get advice if there are infection signs, diabetes/circulation problems, or the blister is large, painful, or unexplained.",
    ),
    VisualCategory(
        "bite",
        "Insect bite or sting",
        "a close-up medical photo of an insect bite, sting, small swollen bump, or localized skin reaction",
        "The image may show a small raised area with localized redness or swelling.",
        "Wash the area, use a cold pack, and avoid scratching.",
        "Seek care for trouble breathing, facial swelling, fever, expanding redness, pus, or a bullseye-like rash after a tick bite.",
    ),
    VisualCategory(
        "eye",
        "Eye redness or irritation",
        "a close-up medical photo of a red irritated eye, conjunctivitis, or eye inflammation",
        "The image may show redness or irritation around the eye.",
        "Avoid rubbing the eye and do not share towels or eye makeup.",
        "Seek urgent care for eye pain, vision changes, chemical exposure, injury, light sensitivity, or contact lens-related redness.",
    ),
    VisualCategory(
        "acne",
        "Acne-like bump or pimple",
        "a close-up medical photo of acne, pimples, clogged pores, or inflamed bumps on skin",
        "The image may show small inflamed bumps or clogged pores.",
        "Avoid squeezing, wash gently, and use non-irritating skin products.",
        "Get advice if lesions are painful, rapidly worsening, leaving scars, or accompanied by fever or spreading redness.",
    ),
    VisualCategory(
        "unclear",
        "Unclear or not enough visual evidence",
        "a blurry unclear photo where the medical skin finding is hard to identify",
        "The image may not contain enough clear visual detail for a useful comparison.",
        "Try a sharper image in good light with the affected area centered, while avoiding identifying features.",
        "Use symptoms and context, not the image alone, to decide whether professional care is needed.",
    ),
]


RED_FLAGS = [
    "trouble breathing, chest pain, fainting, or confusion",
    "rapidly spreading redness, red streaks, pus, or fever",
    "severe pain, numbness, deformity, or inability to move/bear weight",
    "eye pain, vision changes, chemical exposure, or significant eye injury",
    "large/deep wounds, uncontrolled bleeding, animal/human bites, or burns on face/hands/genitals",
]


@functools.lru_cache(maxsize=1)
def load_model():
    processor = CLIPProcessor.from_pretrained(MODEL_ID)
    model = CLIPModel.from_pretrained(MODEL_ID)
    model.eval()
    return processor, model


def classify_image(image: Image.Image):
    processor, model = load_model()
    image = image.convert("RGB")
    prompts = [category.prompt for category in CATEGORIES]
    inputs = processor(text=prompts, images=image, return_tensors="pt", padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        scores = outputs.logits_per_image[0].softmax(dim=0).tolist()

    ranked = sorted(
        zip(CATEGORIES, scores),
        key=lambda item: item[1],
        reverse=True,
    )
    return ranked


def confidence_note(ranked):
    top_score = ranked[0][1]
    second_score = ranked[1][1]
    if top_score < 0.22:
        return "Low confidence: the image does not strongly match any label in this demo set."
    if top_score - second_score < 0.08:
        return "Mixed signal: the top matches are close together, so the explanation should preserve uncertainty."
    return "Clearer visual match within this limited label set, but still not a diagnosis."


def format_results(ranked, style, context):
    top = ranked[:3]
    lines = [
        "## Visual comparison results",
        "",
        "**Important:** this is a research demo, not medical diagnosis. CLIP was not trained as a clinical diagnostic model. Scores are relative visual-match scores for the labels this app provides, not true probabilities.",
        "",
        f"**Uncertainty note:** {confidence_note(ranked)}",
        "",
        "| Possible visual category | Match score | Why it may fit |",
        "| --- | ---: | --- |",
    ]

    for category, score in top:
        lines.append(f"| {category.label} | {score:.2f} | {category.why} |")

    if context.strip():
        lines.extend(
            [
                "",
                "## User context to consider",
                context.strip(),
            ]
        )

    if style == "Simple advice":
        primary = top[0][0]
        lines.extend(
            [
                "",
                "## Plain-language next step",
                f"The closest visual match is **{primary.label.lower()}**, but the image alone is not enough to know what is happening.",
                "",
                f"**Try now:** {primary.home_care}",
                "",
                f"**Do not wait:** {primary.watch_for}",
            ]
        )
    elif style == "Balanced explanation":
        lines.append("")
        lines.append("## What to do with this information")
        for category, _score in top:
            lines.extend(
                [
                    f"**{category.label}:** {category.why}",
                    f"- Care idea: {category.home_care}",
                    f"- Watch for: {category.watch_for}",
                    "",
                ]
            )
    else:
        lines.extend(
            [
                "",
                "## Cautious / uncertainty-first version",
                "Several visible conditions can look similar in a photo. A useful explanation should say what the image resembles, why that comparison is uncertain, and what would change the risk level.",
                "",
            ]
        )
        for category, score in top:
            lines.append(f"- **{category.label} ({score:.2f})**: {category.why} The main safety question is: {category.watch_for}")

    lines.extend(
        [
            "",
            "## Always escalate for",
            *[f"- {item}" for item in RED_FLAGS],
        ]
    )

    return "\n".join(lines)


def analyze(image, explanation_style, symptom_context):
    if image is None:
        return "Upload an image first."

    ranked = classify_image(image)
    return format_results(ranked, explanation_style, symptom_context or "")


EXAMPLES = [
    "Mild redness on arm, started this morning, itchy but not painful.",
    "Twisted ankle yesterday; swelling is worse today and walking hurts.",
    "Small cut from kitchen knife; bleeding stopped after pressure.",
]


with gr.Blocks(title="Health Image Explainer Research Demo") as demo:
    gr.Markdown(
        """
# Health Image Explainer Research Demo
Upload a visible symptom or injury image and compare how different explanation styles handle uncertainty.
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Image",
                type="pil",
                sources=["upload", "webcam", "clipboard"],
            )
            explanation_style = gr.Radio(
                ["Simple advice", "Balanced explanation", "Cautious uncertainty"],
                value="Balanced explanation",
                label="Explanation style",
            )
            symptom_context = gr.Textbox(
                label="Optional symptom context",
                placeholder="What happened? How long has it been there? Pain, fever, spreading, injury, itchiness?",
                lines=4,
            )
            analyze_button = gr.Button("Analyze", variant="primary")

        with gr.Column(scale=2):
            output = gr.Markdown(label="Result")

    gr.Examples(
        examples=[[example] for example in EXAMPLES],
        inputs=[symptom_context],
    )

    analyze_button.click(
        fn=analyze,
        inputs=[image_input, explanation_style, symptom_context],
        outputs=output,
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch()