File size: 1,640 Bytes
51bf14d
c0ca15d
 
 
 
51bf14d
c0ca15d
 
 
 
 
 
 
 
51bf14d
c0ca15d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
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()