import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image

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

# Define the core prediction function
def predict_flower(image):
    # Preprocess image
    image = image.resize((150, 150))  # Resize the image to 150x150
    image = image.convert('RGB')  # Ensure image has 3 channels
    image = np.array(image) / 255.0  # Normalize the image to [0, 1]
    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 Flower classes
    class_names = ['daisy', 'dandelion', 'rose', 'sunflower', 'tulip']
    probabilities_dict = {flower_class: round(float(probability), 2) for flower_class, probability in zip(class_names, probabilities.numpy()[0])}
    
    return probabilities_dict

# Streamlit interface
st.title("Der Bluemenerkenner")
st.write("Welche Blume wächst in ihrem Garten?")

# Upload image
uploaded_image = st.file_uploader("Lade dein Bild hoch...", type=["jpg", "png"])

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption='Hochgeladenes Bild.', use_column_width=True)
    st.write("")
    st.write("Identifizieren...")
    
    predictions = predict_flower(image)
    
    # Log the predictions
    st.write("Vorhersagen für jede Klasse:")
    for flower_class, probability in predictions.items():
        st.write(f"{flower_class}: {probability}")
    
    # Find the flower with the highest probability
    highest_probability_flower = max(predictions, key=predictions.get)
    
    
    # Display the flower with the highest probability
    st.write(f"Es ist eine: **{highest_probability_flower}**")