import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt

# Load the trained model
model_path = "skin-model_transferlearning5.keras"
model = tf.keras.models.load_model(model_path)

# Define the core prediction function
def predict_skin(image):
    # Preprocess image
    image = image.resize((450, 450))  # Resize the image to 450x450
    image = image.convert('RGB')  # Ensure image has 3 channels
    image = np.array(image)
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    
    # Predict
    prediction = model.predict(image)
    
    # Apply softmax to get probabilities for each class
    probabilities = tf.nn.softmax(prediction, axis=1)
    
    # Map probabilities to Skin condition classes
    class_names = ['Benign', 'Malicious']
    probabilities_dict = {skin_class: round(float(probability), 2) for skin_class, probability in zip(class_names, probabilities.numpy()[0])}
    
    return probabilities_dict

# Streamlit interface
st.title("Skin disease detection")
st.write(" The Skin Disease Detector is a prototype and should not be used as the sole method for diagnosing skin diseases. For accurate diagnosis and treatment, please consult a trained medical professional. This tool is intended for preliminary screening purposes only and should not replace professional medical advice.")
st.write("")
st.write("Accurate classification works best if you provide a clear and close-up image with natural lighting.")
st.write("")

# Upload image
uploaded_image = st.file_uploader("Choose an image", type=["jpg", "png"])

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Classifying...")
    
    predictions = predict_skin(image)
    
    # Display predictions as a DataFrame
    st.write("### Probabilities")
    df = pd.DataFrame(predictions.items(), columns=["Condition", "Probability"])
    st.dataframe(df)
    
    # Display predictions as a bar chart
    st.write("### Prediction Chart")
    fig, ax = plt.subplots()
    ax.barh(df["Condition"], df["Probability"], color='skyblue')
    ax.set_xlim(0, 1)
    ax.set_xlabel('Probability')
    ax.set_title('Prediction Probabilities')
    st.pyplot(fig)