Update app.py

app.py

@@ -1,38 +1,39 @@
-
 import gradio as gr
 import tensorflow as tf
-from PIL import Image
 import numpy as np
+from PIL import Image
 
-labels = ['Haunter', 'Gengar', 'Ditto', 'Vulpix']
-
-def predict_pokemon_type(uploaded_file):
-    if uploaded_file is None:
-        return "No file uploaded.", None, "No prediction"
-
-    model = tf.keras.models.load_model('pokemon-model_2_transferlearning.keras')
-
-    # Load the image from the file path
-    with Image.open(uploaded_file) as img:
-        img = img.resize((150, 150))
-        img_array = np.array(img)
-
-        prediction = model.predict(np.expand_dims(img_array, axis=0))
-        confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
-        
-        # Identify the most confident prediction
-        confidences = {labels[i]: np.round(float(prediction[0][i]), 2) for i in range(len(labels))}
-
-        return img, confidences
 
-# Define the Gradio interface
+model_path = "pokemon-model_2_transferlearning.keras"
+model = tf.keras.models.load_model(model_path)
+
+def predict_pokemon(image):
+    # Preprocess image
+    print(type(image))
+    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL image
+    image = image.resize((150, 150)) #resize the image to 150x150
+    image = np.array(image)
+    image = np.expand_dims(image, axis=0) # same as image[None, ...]
+    
+    
+    prediction = model.predict(image)
+    
+    # Apply softmax to get probabilities for each class
+    prediction = tf.nn.softmax(prediction)
+    
+    # Create a dictionary with the probabilities for each Pokemon
+    evee = np.round(float(prediction[0][0]), 2)
+    farfetched = np.round(float(prediction[0][1]), 2)
+    graveler = np.round(float(prediction[0][2]), 2)
+    venonta = np.round(float(prediction[0][3]), 2)
+    
+    return {'Evee': evee, 'Farfetched': farfetched, 'Graveler': graveler, 'Venonta': venonta}
+
+
+input_image = gr.Image()
 iface = gr.Interface(
-    fn=predict_pokemon_type,  # Function to process the input
-    inputs=gr.File(label="Upload File"),  # File upload widget
-    outputs=["image", "text"],  # Output types for image and text
-    title="Pokemon Classifier",  # Title of the interface
-    description="Upload a picture of a Pokemon (preferably Ditto, Gengar, Haunting, or Vulpix) to see its type and confidence level."
-)
-
-
-iface.launch()
+    fn=predict_pokemon,
+    inputs=input_image, 
+    outputs=gr.Label(),
+    description="A simple mlp classification model for image classification using the mnist dataset.")
+iface.launch(share=True)