Update app.py

app.py

@@ -2,6 +2,8 @@ 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 = "pokemon-model_transferlearning1.keras"
@@ -11,7 +13,7 @@ model = tf.keras.models.load_model(model_path)
 def predict_pokemon(image):
     # Preprocess image
     image = image.resize((150, 150))  # Resize the image to 150x150
-    image = image.convert('RGB')  # Ensure the image is in RGB format
+    image = image.convert('RGB')  # Ensure image has 3 channels
     image = np.array(image)
     image = np.expand_dims(image, axis=0)  # Add batch dimension
     
@@ -29,7 +31,7 @@ def predict_pokemon(image):
 
 # Streamlit interface
 st.title("Pokemon Classifier")
-st.write("Eine KI die Pokemons identifiziert :) Viel Spass")
+st.write("A simple MLP classification model for image classification using a pretrained model.")
 
 # Upload image
 uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png"])
@@ -42,11 +44,22 @@ if uploaded_image is not None:
     
     predictions = predict_pokemon(image)
     
-    st.write(predictions)
+    # Display predictions as a DataFrame
+    st.write("### Prediction Probabilities")
+    df = pd.DataFrame(predictions.items(), columns=["Pokemon", "Probability"])
+    st.dataframe(df)
+    
+    # Display predictions as a bar chart
+    st.write("### Prediction Chart")
+    fig, ax = plt.subplots()
+    ax.barh(df["Pokemon"], df["Probability"], color='skyblue')
+    ax.set_xlim(0, 1)
+    ax.set_xlabel('Probability')
+    ax.set_title('Prediction Probabilities')
+    st.pyplot(fig)
 
 # Example images
 st.sidebar.title("Examples")
 example_images = ["pokemon/train/chansey/00000000.png", "pokemon/train/growlithe/00000000.png", "pokemon/train/lapras/00000000.png"]
 for example_image in example_images:
     st.sidebar.image(example_image, use_column_width=True)
-