Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
from PIL import Image | |
model_path = "transfer_learning_pk.keras" | |
model = tf.keras.models.load_model(model_path) | |
# Define the core prediction function | |
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 28x28 and converts it to gray scale | |
image = np.array(image) | |
image = np.expand_dims(image, axis=0) # same as image[None, ...] | |
# Predict | |
prediction = model.predict(image) | |
# No need to apply sigmoid, as the output layer already uses softmax | |
# Convert the probabilities to rounded values | |
prediction = np.round(prediction, 2) | |
# Separate the probabilities for each class | |
p_gyarados = prediction[0][0] # Probability for class 'articuno' | |
p_metapod = prediction[0][1] # Probability for class 'moltres' | |
p_ponyta = prediction[0][2] # Probability for class 'zapdos' | |
return {'gyarados': p_gyarados, 'metapod': p_metapod, 'ponyta': p_ponyta} | |
# Create the Gradio interface | |
input_image = gr.Image() | |
iface = gr.Interface( | |
fn=predict_pokemon, | |
inputs=input_image, | |
outputs=gr.Label(), | |
examples=["img/gyarados1.jpg", "img/gyarados2.jpg", "img/gyarados3.jpg", "img/metapod1.jpg", "img/metapod2.jpg", "img/metapod3.jpg", "img/ponyta1.jpg", "img/ponyta2.jpg", "img/ponyta3.jpg"], | |
description="Pokemon") | |
iface.launch() | |