Spaces:
Sleeping
Sleeping
import gradio as gr | |
from PIL import Image | |
import numpy as np | |
from tensorflow.keras.preprocessing import image as keras_image | |
from tensorflow.keras.applications.resnet50 import preprocess_input as resnet_preprocess_input | |
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input as mobilenet_preprocess_input | |
from tensorflow.keras.models import load_model | |
# Load your trained models | |
resnet_model = load_model('/Users/beharademi/Documents/ZHAW/6. Sem LOKAL/KI/module exam/ademibeh/resnet_best_model.keras') # Update path | |
mobilenet_model = load_model('/Users/beharademi/Documents/ZHAW/6. Sem LOKAL/KI/module exam/ademibeh/mobilenet_best_model.keras') # Update path | |
def predict_comic_character(img, model_type): | |
img = Image.fromarray(img.astype('uint8'), 'RGB') | |
img = img.resize((224, 224)) # Resize the image to fit model input | |
img_array = keras_image.img_to_array(img) | |
img_array = np.expand_dims(img_array, axis=0) | |
if model_type == 'ResNet50': | |
img_array = resnet_preprocess_input(img_array) | |
prediction = resnet_model.predict(img_array) | |
elif model_type == 'MobileNetV2': | |
img_array = mobilenet_preprocess_input(img_array) | |
prediction = mobilenet_model.predict(img_array) | |
else: | |
return {"error": "Invalid model type selected"} | |
classes = ['Superman', 'Batman', 'WonderWoman', 'Riddler', 'Spider-Man', 'Iron-Man', | |
'Hulk', 'The Joker', 'Magneto', 'Wolverine', 'Deadpool', 'Catwoman'] | |
return {classes[i]: float(prediction[0][i]) for i in range(len(classes))} | |
# Define the Gradio interface | |
interface = gr.Interface( | |
fn=predict_comic_character, | |
inputs="image", | |
outputs="label", | |
title="Comic Character Classifier", | |
description="Upload an image of a comic character and the classifier will predict the character.", | |
) | |
# Launch the interface | |
interface.launch() | |