# app.py

import gradio as gr
from PIL import Image
import torch
import pandas as pd
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load model
model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
label_map = model.config.id2label

# Match labels exactly from model
condition_info = {
    "actinic_keratoses": "Dry, rough patch – sometimes early sign of skin cancer.",
    "basal_cell_carcinoma": "Slow-growing skin cancer. Common but treatable.",
    "benign_keratosis-like_lesions": "Non-cancerous growth. Like age spots or warts.",
    "dermatofibroma": "Small, firm bump. Usually harmless.",
    "melanocytic_nevi": "Just a mole. Normal unless changing.",
    "melanoma": "Dangerous skin cancer. Needs fast treatment.",
    "vascular_lesions": "Red or purple patches from blood vessels."
}

# AI prediction logic
def classify_skin(image: Image.Image):
    if image is None:
        return pd.DataFrame(), "Please upload or take a photo."

    image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=1)[0]

    threshold = 0.40
    data = []
    likely_conditions = []

    for idx, prob in enumerate(probs):
        label = label_map[idx]  # this is the exact model label
        conf = prob.item()
        status = "✅ Positive" if conf > threshold else "❌ Negative"
        desc = condition_info.get(label, "No description available.")
        data.append({
            "Condition": label.replace("_", " ").capitalize(),
            "Confidence (%)": f"{conf*100:.2f}",
            "Status": status,
            "What it means": desc
        })
        if conf > threshold:
            likely_conditions.append(label.replace("_", " ").capitalize())

    df = pd.DataFrame(data)
    summary_text = (
        "🧾 **Summary:** " +
        (", ".join(likely_conditions) if likely_conditions else "No major concern seen by AI.") +
        "\n\n📢 Please check with a real doctor for correct diagnosis."
    )

    return df, summary_text

# Gradio UI
demo = gr.Interface(
    fn=classify_skin,
    inputs=gr.Image(type="pil", label="📷 Upload or Capture Skin Image"),
    outputs=[
        gr.Dataframe(headers=["Condition", "Confidence (%)", "Status", "What it means"]),
        gr.Markdown()
    ],
    title="AI Skin Condition Classifier",
    description="Upload a photo of a skin issue. The AI will check 7 common conditions and suggest what's likely. For support only — not a diagnosis."
)

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