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

# Cache the model and feature extractor
@st.cache_resource
def load_model():
    model_name = "syaha/skin_cancer_detection_model"
    model = AutoModelForImageClassification.from_pretrained(model_name)
    feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
    return model, feature_extractor

model, feature_extractor = load_model()

# App title and instructions
st.title("Skin Cancer Detection App")
st.write("""
Upload an image of the affected skin area, and this app will classify it based on pre-trained skin lesion types.
""")

# File uploader for user input
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    try:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Preprocess the image
        inputs = feature_extractor(images=image, return_tensors="pt")
        
        # Perform inference
        with torch.no_grad():
            outputs = model(**inputs)
            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
            predicted_class = probabilities.argmax().item()
            confidence = probabilities.max().item()

        # Map predicted class to human-readable labels
        labels = {
            0: "Melanocytic nevi",
            1: "Melanoma",
            2: "Benign keratosis-like lesions",
            3: "Basal cell carcinoma",
            4: "Actinic keratoses",
            5: "Vascular lesions",
            6: "Dermatofibroma"
        }
        diagnosis = labels.get(predicted_class, "Unknown")

        # Display results
        st.subheader("Results")
        st.write(f"**Diagnosis:** {diagnosis}")
        st.write(f"**Confidence Level:** {confidence:.2%}")

        # Provide an option to download the report
        report = f"Diagnosis: {diagnosis}\nConfidence Level: {confidence:.2%}"
        st.download_button(
            label="Download Report",
            data=report,
            file_name="report.txt",
            mime="text/plain",
        )
    except Exception as e:
        st.error(f"An error occurred while processing the image: {e}")