# app.py

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

# Set up Streamlit page
st.set_page_config(page_title="Skin AI Classifier", layout="centered")
st.title("AI-Based Skin Condition Classifier")
st.markdown("Upload a clear skin photo. The AI will suggest the likely skin condition.")

# Upload image
uploaded_file = st.file_uploader("Upload a skin photo (JPG, PNG)", type=["jpg", "jpeg", "png"])

# Load pre-trained model from Hugging Face (Anwarkh1)
@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()

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

    # Preprocess and predict
    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 predictions
    top_probs, top_indices = torch.topk(probs, k=3)
    class_names = model.config.id2label

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

    st.caption("Note: This AI tool is for support only. Always consult a certified dermatologist.")