Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
from PIL import Image | |
# Load the model | |
model_path = "model_2_familyguy.keras" | |
model = tf.keras.models.load_model(model_path) | |
# Define the core prediction function | |
def predict_familyguy(image): | |
# Preprocess 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) / 255.0 # Normalize the image | |
image = np.expand_dims(image, axis=0) # Add batch dimension | |
# Predict | |
prediction = model.predict(image) | |
# Convert the probabilities to rounded values | |
prediction = np.round(prediction, 2) | |
# Separate the probabilities for each class | |
p_brian = prediction[0][0] # Probability for class 'Brain Griffin' | |
p_lois = prediction[0][1] # Probability for class 'Lois Griffin' | |
p_peter = prediction[0][2] # Probability for class 'Peter Griffin' | |
p_stewie = prediction[0][3] # Probability for class 'Stewie Griffin' | |
return {'brian': p_brian, 'lois': p_lois, 'peter': p_peter, 'stewie': p_stewie} | |
# Create the Gradio interface | |
input_image = gr.Image() | |
iface = gr.Interface( | |
fn=predict_familyguy, | |
inputs=input_image, | |
outputs=gr.Label(num_top_classes=4), # This will display the top 4 class labels | |
examples=["images/Brain1.png", "images/Brain2.jpg", "images/Brain3.jpg", | |
"images/Lois1.png", "images/Lois2.png", "images/Lois3.png", | |
"images/Peter1.jpg", "images/Peter2.jpg", "images/Peter3.jpg", | |
"images/Stewie1.jpg", "images/Stewie2.jpg", "images/Stewie3.jpg"], | |
description="Upload an image to classify it as Brian, Lois, Peter, or Stewie." | |
) | |
iface.launch() | |