pokemon / app.py
hasvio01's picture
Update app.py
2c8878b verified
import tensorflow as tf
import gradio as gr
import numpy as np
import os
from PIL import Image
print(tf.__version__)
print(f"Current working directory: {os.getcwd()}")
model_path = 'transferlearning_pokemon.h5'
# Check if the model exists
if os.path.exists(model_path):
print(f"Model found at {model_path}")
try:
# Load the trained model
model = tf.keras.models.load_model(model_path)
print("Model loaded successfully.")
except Exception as e:
print(f"Error loading model: {e}")
else:
print(f"Model not found at {model_path}. Please check the path.")
# Define class names (make sure this matches the classes used during training)
class_names = ['Machamp', 'Raichu', 'Vulpix']
# Define the prediction function
def predict(image):
image = Image.fromarray(image.astype('uint8')) # Convert numpy array to PIL image
image = image.resize((150, 150)) #resize the image to 28x28 and converts it to gray scale
image = np.array(image)
image = np.expand_dims(image, axis=0) # same as image[None, ...]
predictions = model.predict(image)
predicted_class = np.argmax(predictions, axis=1)[0]
confidence = np.max(predictions)
return {class_names[predicted_class]: float(confidence)}
# Create a Gradio interfac
input_image = gr.Image()
output_text = gr.Textbox(label="Predicted Value")
interface = gr.Interface(fn=predict,
inputs=input_image,
outputs=gr.Label(),
examples=["00000000.jpg", "00000001.jpg", "00000010.png", "00000017.jpg", "00000021.jpg", "00000067.jpg"],
description="A simple mlp classification model for image classification using the mnist dataset.")
# Launch the Gradio interface
interface.launch()