Spaces:
Runtime error
Runtime error
import gradio as gr | |
import cv2 | |
import matplotlib.pyplot as plt | |
import numpy as np | |
from openvino.runtime import Core | |
##### | |
#Load pretrained model | |
##### | |
ie = Core() | |
model_path = "./model/v3-small_224_1.0_float.xml" | |
model = ie.read_model(model=model_path) | |
compiled_model = ie.compile_model(model=model, device_name="CPU") | |
output_layer = compiled_model.output(0) | |
##### | |
#Inference | |
##### | |
def predict(img: np.ndarray) -> str: | |
# input: numpy array of image in RGB (see defaults for https://www.gradio.app/docs/#image) | |
print(f'initial image shape: {img.shape}') | |
# The MobileNet model expects images in RGB format. | |
# Resize to MobileNet image shape. | |
input_image = cv2.resize(src=img, dsize=(224, 224)) | |
print(f'resized: {input_image.shape}') | |
# Reshape to model input shape. | |
input_image = np.expand_dims(input_image, 0) | |
print(f'final shape: {input_image.shape}') | |
# Get inference result | |
result_infer = compiled_model([input_image])[output_layer] | |
result_index = np.argmax(result_infer) | |
# Convert the inference result to a class name. | |
imagenet_classes = open("./model/imagenet_2012.txt").read().splitlines() | |
# The model description states that for this model, class 0 is a background. | |
# Therefore, a background must be added at the beginning of imagenet_classes. | |
imagenet_classes = ['background'] + imagenet_classes | |
best_class = imagenet_classes[result_index] | |
# clean up | |
best_class = best_class.partition(' ')[2] | |
# TODO: get n best results with corresponding probabilities? | |
return best_class | |
##### | |
#Gradio Setup | |
##### | |
title = "Image classification" | |
description = "Image classification with OpenVino model trained on ImageNet" | |
examples = ['dog.jpg'] | |
interpretation='default' | |
enable_queue=True | |
gr.Interface( | |
fn=predict, | |
inputs=gr.inputs.Image(), | |
outputs=gr.outputs.Label(num_top_classes=1), | |
title=title, | |
description=description, | |
examples=examples, | |
interpretation=interpretation, | |
enable_queue=enable_queue | |
).launch() |