|
import os |
|
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' |
|
|
|
import gradio as gr |
|
import tensorflow as tf |
|
from PIL import Image |
|
import numpy as np |
|
|
|
|
|
model_path = "pokemon_classifier_model.keras" |
|
|
|
|
|
model = tf.keras.models.load_model(model_path) |
|
|
|
|
|
labels = ['Abra', 'Cloyster', 'Dodrio'] |
|
|
|
|
|
def predict(image): |
|
try: |
|
|
|
image = image.resize((150, 150)) |
|
image = np.array(image) / 255.0 |
|
image = np.expand_dims(image, axis=0) |
|
|
|
|
|
predictions = model.predict(image) |
|
confidences = {labels[i]: float(predictions[0][i]) for i in range(len(labels))} |
|
|
|
return confidences |
|
except Exception as e: |
|
return str(e) |
|
|
|
|
|
examples = [ |
|
["examples/abra.png"], |
|
["examples/cloyster.png"], |
|
["examples/dodrio.png"] |
|
] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
inputs=gr.Image(type="pil"), |
|
outputs=gr.Label(), |
|
examples=examples, |
|
description="Pokémon Classifier" |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |
|
|