Update app.py

app.py

@@ -27,7 +27,10 @@ def predict_pokemon(image):
     class_names = ['Pikachu', 'Raichu', 'Snorlax']
     probabilities_dict = {pokemon_class: round(float(probability), 2) for pokemon_class, probability in zip(class_names, probabilities.numpy()[0])}
     
-    return probabilities_dict
+    # Get the predicted Pokemon
+    predicted_pokemon = max(probabilities_dict, key=probabilities_dict.get)
+    
+    return predicted_pokemon, probabilities_dict
 
 # Streamlit interface
 st.title("Pokemon Classification Tool")
@@ -42,18 +45,15 @@ if uploaded_image is not None:
     st.write("")
     st.write("Classifying...")
     
-    predictions = predict_pokemon(image)
+    predicted_pokemon, predictions = predict_pokemon(image)
     
-    # Display predictions as a DataFrame
-    st.write("### Prediction Probabilities")
-    df = pd.DataFrame(predictions.items(), columns=["Pokemon", "Probability"])
-    st.dataframe(df)
+    # Display predicted Pokemon name in bold and large font
+    st.write(f"### **Predicted Pokemon: {predicted_pokemon.upper()}**")
     
-    # Display predictions as a bar chart
+    # Display predictions as a pie 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.pie(predictions.values(), labels=predictions.keys(), autopct='%1.1f%%', startangle=90)
     ax.set_title('Prediction Probabilities')
-    st.pyplot(fig)
+    ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
+    st.pyplot(fig)