LuisDarioHinojosa
fixed description
520adc6
import os
import numpy as np
from timeit import default_timer as timer
import cv2 as cv
import gradio as gr
from model_instance_function import get_pretrained_dog_emotion_classifier
# normalize function
def image_preprocessing(img):
img = np.array(img)
img = cv.resize(img,(224,224))
img = img.reshape(1,224,224,3)
return img / 255.0
# instance the model
model = get_pretrained_dog_emotion_classifier()
# gradio predict function
def predict(img):
# class to map the indices to the classes
class_2_index = {0: 'happy', 1: 'sad'}
# measure execution time
start_time = timer()
# preprocess the image
img = image_preprocessing(img)
# make a prediction (prob of sad dog)
pred_probability = model.predict(img)[0]
# convert to an index
pred_index = 1 if pred_probability > 0.5 else 0
# label
pred_label = class_2_index[pred_index]
end_time = timer()
total_time = end_time - start_time
return pred_probability, pred_label,round(total_time,5)
title = "Dog Emotions Vision Classifier"
description = "A vision classifier that distinguishes between sad and happy dogs."
article = "The model was trained in the [Dogs Emotions Dataset](https://huggingface.co/datasets/Q-b1t/Dogs_Emotions_Dataset) using the pretrained convolutional blocks of the VGG16 architecture and a custom classifier. For more information regarding the training, refer to this [colab notebook](https://colab.research.google.com/drive/1QqjLFsNV_8N1xr29BVwn4QVs_VH6lXmV?usp=sharing)."
example_list = [["examples/" + example] for example in os.listdir("examples")]
demo = gr.Interface(
fn = predict,
inputs = gr.Image(type = "pil"),
outputs = [gr.Number(label = "Probability of a sad dog"),gr.Textbox(max_lines = 2,label = "Most likely class"),gr.Number(label = "Prediction time (s)")],
examples = example_list,
title = title,
description = description,
article = article
)
demo.launch()