Sa-m's picture
Update app.py
03c6e36 verified
import gradio as gr
import tensorflow as tf
import numpy as np
import os
# Configuration
HEIGHT, WIDTH = 224, 224
NUM_CLASSES = 6
LABELS = ["McDonalds", "Burger King", "Subway", "Starbucks", "KFC", "Other"]
from tensorflow_addons.metrics import F1Score
from keras.utils import custom_object_scope
with custom_object_scope({'Addons>F1Score': F1Score}):
model = tf.keras.models.load_model('best_model.h5')
def classify_image(inp):
# Resize & preprocess
inp = tf.image.resize(inp, [HEIGHT, WIDTH])
inp = tf.cast(inp, tf.float32)
inp = tf.keras.applications.nasnet.preprocess_input(inp)
inp = tf.expand_dims(inp, axis=0)
# Predict
prediction = model.predict(inp)[0]
return {LABELS[i]: float(f"{prediction[i]:.6f}") for i in range(NUM_CLASSES)}
example_list = [
["Examples/Untitled.png"],
["Examples/Untitled2.png"],
["Examples/Untitled3.png"],
["Examples/Untitled5.png"]
]
iface = gr.Interface(
fn=classify_image,
inputs=gr.Image(
label="Input Image",
sources="upload",
type="numpy",
height=HEIGHT,
width=WIDTH
),
outputs=gr.Label(num_top_classes=4),
title="Brand Logo Detection",
examples=example_list
)
if __name__ == "__main__":
iface.launch(debug=False, share=True)