Comic-Charakter / app.py
ademibeh's picture
Upload app.py
a21c3cc verified
raw
history blame
1.88 kB
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()