# app.py

import streamlit as st
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

# App title and instructions
st.set_page_config(page_title="Skin Condition Classifier", layout="centered")
st.title("🧠 AI Skin Condition Classifier")
st.markdown("Upload a **clear photo** of the skin condition to receive AI-powered predictions.")

# Image uploader
uploaded_file = st.file_uploader("📷 Upload a skin image", type=["jpg", "jpeg", "png"])

# Load the pre-trained model
@st.cache_resource
def load_model():
    model_name = "Anwarkh1/Skin_Cancer-Image_Classification"
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = AutoModelForImageClassification.from_pretrained(model_name)
    return processor, model

processor, model = load_model()

# Handle image upload and prediction
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    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]

    # Top 3 predictions
    top_probs, top_indices = torch.topk(probs, k=3)
    class_labels = model.config.id2label

    st.subheader("🧾 Prediction Results")
    for idx, prob in zip(top_indices, top_probs):
        label = class_labels[idx.item()]
        st.write(f"**{label}** – {prob.item() * 100:.2f}%")

    st.info("🔍 Note: This tool is for supportive use only. Please consult a dermatologist for a medical diagnosis.")