Update app.py

app.py

@@ -11,6 +11,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 = np.array(image)
     image = np.expand_dims(image, axis=0)  # Add batch dimension
     
@@ -22,13 +23,13 @@ def predict_pokemon(image):
     
     # Map probabilities to Pokemon classes
     class_names = ['Chansey', 'Growlithe', 'Lapras']
-    probabilities_dict = {pokemon_class: round(float(probability), 2) for pokemon_class, probability in zip(pokemon_classes, probabilities.numpy()[0])}
+    probabilities_dict = {pokemon_class: round(float(probability), 2) for pokemon_class, probability in zip(class_names, probabilities.numpy()[0])}
     
     return probabilities_dict
 
 # Streamlit interface
 st.title("Pokemon Classifier")
-st.write("A simple MLP classification model for image classification using a pretrained model.")
+st.write("Eine KI die Pokemons identifiziert :) Viel Spass")
 
 # Upload image
 uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png"])
@@ -48,3 +49,4 @@ 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)
+