rbarman's picture
type hints + simplify code if input is numpy array
6bad718
raw
history blame
1.82 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: np.ndarray) -> str:
# The MobileNet model expects images in RGB format.
image = cv2.cvtColor(img, code=cv2.COLOR_BGR2RGB)
# Resize to MobileNet image shape.
input_image = cv2.resize(src=image, dsize=(224, 224))
# Reshape to model input shape.
input_image = np.expand_dims(input_image, 0)
# 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]
# 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()