import gradio as gr
from openvino.inference_engine import IECore
import cv2
import numpy as np

# Load the OpenVINO model
ie = IECore()
net = ie.read_network(model='openvino/best.xml', weights='openvino/best.bin')
exec_net = ie.load_network(network=net, device_name="CPU")

# Define a function for prediction
def predict(image):
    # Preprocess the image
    input_blob = next(iter(net.input_info))
    input_shape = net.input_info[input_blob].input_data.shape
    image_resized = cv2.resize(image, (input_shape[3], input_shape[2]))
    image_transposed = image_resized.transpose((2, 0, 1))
    image_batch = np.expand_dims(image_transposed, axis=0)

    # Perform inference
    result = exec_net.infer(inputs={input_blob: image_batch})
    output_blob = next(iter(net.outputs))
    detections = result[output_blob]

    # Post-process the results
    boxes = []
    for detection in detections[0][0]:
        confidence = detection[2]
        if confidence > 0.5:  # Threshold
            xmin = int(detection[3] * image.shape[1])
            ymin = int(detection[4] * image.shape[0])
            xmax = int(detection[5] * image.shape[1])
            ymax = int(detection[6] * image.shape[0])
            boxes.append((xmin, ymin, xmax, ymax, confidence))

    # Draw bounding boxes on the image
    for (xmin, ymin, xmax, ymax, confidence) in boxes:
        cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)

    return image

# Create a Gradio interface
iface = gr.Interface(fn=predict, inputs=gr.Image(type="numpy"), outputs=gr.Image(type="numpy"))

# Launch the Gradio app
iface.launch()