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()