rbarman's picture
try openvino model
c0ca15d
raw
history blame
1.64 kB
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):
#img = PILImage.create(img)
#pred,pred_idx,probs = learn.predict(img)
#return {labels[i]: float(probs[i]) for i in range(len(labels))}
# TODO: get n best results with corresponding probabilities?
# 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
return imagenet_classes[result_index]
#####
#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(shape=(512, 512)),
outputs=gr.outputs.Label(num_top_classes=1),
title=title,
description=description,
examples=examples,
interpretation=interpretation,
enable_queue=enable_queue
).launch()