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

# Set the title of the application
st.title("Dermavision")
st.write(
    "Upload an image of the affected skin area, and the app will classify the disease and provide analysis."
)

# Cache model and processor loading
@st.cache_resource
def load_model():
    repo_name = "Jayanth2002/dinov2-base-finetuned-SkinDisease"
    processor = AutoImageProcessor.from_pretrained(repo_name)
    model = AutoModelForImageClassification.from_pretrained(repo_name)
    return model, processor

model, processor = load_model()

# Define the class names
class_names = [
    'Basal Cell Carcinoma', 'Darier_s Disease', 'Epidermolysis Bullosa Pruriginosa',
    'Hailey-Hailey Disease', 'Herpes Simplex', 'Impetigo', 'Larva Migrans',
    'Leprosy Borderline', 'Leprosy Lepromatous', 'Leprosy Tuberculoid', 'Lichen Planus',
    'Lupus Erythematosus Chronicus Discoides', 'Melanoma', 'Molluscum Contagiosum',
    'Mycosis Fungoides', 'Neurofibromatosis', 'Papilomatosis Confluentes And Reticulate',
    'Pediculosis Capitis', 'Pityriasis Rosea', 'Porokeratosis Actinic', 'Psoriasis',
    'Tinea Corporis', 'Tinea Nigra', 'Tungiasis', 'actinic keratosis', 'dermatofibroma',
    'nevus', 'pigmented benign keratosis', 'seborrheic keratosis', 'squamous cell carcinoma',
    'vascular lesion'
]

# Define reasons, treatments, and home remedies for each disease
# (This section is omitted for brevity but should remain unchanged from your original code)

# Function to classify the image
def classify_image(image):
    inputs = processor(image.convert("RGB"), return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class_idx = outputs.logits.argmax(-1).item()
        confidence_score = torch.nn.functional.softmax(outputs.logits, dim=-1).max().item()
    predicted_label = class_names[predicted_class_idx]
    return predicted_label, confidence_score

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

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Analyze the image
    with st.spinner("Analyzing the image..."):
        predicted_label, confidence_score = classify_image(image)
        if predicted_label not in disease_analysis:
            st.error("Unable to classify the disease. Please upload a clearer image or consult a dermatologist.")
        else:
            reason = disease_analysis.get(predicted_label, {}).get("reason", "Reason unknown.")
            treatment = disease_analysis.get(predicted_label, {}).get("treatment", "Consult a dermatologist.")
            home_remedy = disease_analysis.get(predicted_label, {}).get("home_remedy", "No specific home remedies available.")

            # Display the results
            st.success("Analysis Complete!")
            st.markdown(f"### **Classification**: {predicted_label}")
            st.markdown(f"**Confidence**: {confidence_score:.2%}")
            st.markdown(f"**Reason**: {reason}")
            st.markdown(f"**Treatment**: {treatment}")
            st.markdown(f"**Home Remedy**: {home_remedy}")
            st.markdown("**Note:** Please consult a doctor for final recommendations and a detailed treatment plan.")

# Optional feedback form
st.markdown("---")
st.header("We Value Your Feedback!")
feedback = st.text_area("Please share your feedback to help us improve:")
if st.button("Submit Feedback"):
    if feedback:
        st.success("Thank you for your feedback!")
    else:
        st.warning("Feedback cannot be empty.")