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