pokemon_v2 / app.py
fahrnphi's picture
Update app.py
ca58299 verified
raw
history blame
1.71 kB
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_fahrnphi_transferlearning.keras"
model = tf.keras.models.load_model(model_path)
# Define the core prediction function
def predict_pokemon(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)
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 Pokemon classes
class_names = ['Charizard', 'Lapras', 'Machamp']
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.")
# Upload image
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png"])
if uploaded_image is not None:
image = Image.open(uploaded_image)
st.image(image, caption='Uploaded Image.', use_column_width=True)
st.write("")
st.write("Classifying...")
predictions = predict_pokemon(image)
# Display predictions as a DataFrame
st.write("### Prediction Probabilities")
df = pd.DataFrame(predictions.items(), columns=["Pokemon", "Probability"])
st.dataframe(df)