ana-bernal's picture
corrected # of classes
9777e8f
import gradio as gr
# For loading files
from joblib import dump, load
# Model hub
import tensorflow_hub as hub
# Neural networks
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from keras.applications.vgg16 import preprocess_input
from huggingface_hub import from_pretrained_keras
# Image processing
import PIL
#------------------------------------------
# Loading model
model = from_pretrained_keras('ana-bernal/keras_dog_breed_eff')
# Reading file with class names, uncomment to import class names
breed_names_norm = []
with open('labels.txt', 'r') as file:
for line in file:
# remove linebreak from a current name
name = line[:-1]
breed_names_norm.append(name)
# Definition of main function
def classify_image(inp):
"""
Returns a dictionnary: predicted_breeds, where the
keys are [1,2,3] for the first, second and third more probable
breed for the dog image. Each value is a dictionnary with keys
['idx', 'name', 'confidence'] and their corresponding values.
Parameters:
img: returned by the function load_img_path
"""
img_array = keras.preprocessing.image.img_to_array(inp)
img_array = tf.expand_dims(img_array, 0) # Creates a batch axis
predictions = model.predict(img_array, verbose=0).flatten()
confidences = {breed_names_norm[i]: float(predictions[i]) for i in range(120)}
return confidences
# --------------------------------------------------
examples = [
['example_images/01_test.jpg'],
['example_images/02_test.jpg'],
['example_images/03_test.jpg'],
['example_images/04_test.jpg'],
['example_images/05_test.jpg'],
['example_images/06_test.jpg'],
]
demo = gr.Interface(fn=classify_image,
inputs=gr.Image(shape=(180, 180)),
outputs=gr.Label(num_top_classes=3),
examples=examples)
if __name__ == "__main__":
demo.launch()