rbarman's picture
bruhhhhh
69d9453
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()