import gradio as gr
import numpy as np
import cv2
from openvino.runtime import Core

# Load the OpenVINO model
ie = Core()
model_path = "best_int8_openvino_model/best.xml" 
model = ie.read_model(model=model_path)
compiled_model = ie.compile_model(model=model, device_name="CPU")

# Get input and output layers
input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)

# Function to process the image and make predictions
def predict(image):
    # Preprocess the image
    image_resized = cv2.resize(image, (input_layer.shape[-2], input_layer.shape[-1]))
    input_image = np.expand_dims(image_resized.transpose(2, 0, 1), 0)  # NHWC to NCHW
    
    # Run inference
    result = compiled_model([input_image])[output_layer]
    
    # Process the result    
    detected_boxes = result[:, :4]
    scores = result[:, 4]
    
    # Draw the boxes on the image
    for box, score in zip(detected_boxes, scores):
        if score[0] > 0.5:
            x_min, y_min, x_max, y_max = box.astype(int)
            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
            cv2.putText(image, f"{score[0]:.2f}", (x_min, y_min-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
    
    return image

# Define the Gradio interface
gr_interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title="ITI107 2024S2 Assignment: YOLOv8 Object Detection",
    description="Upload an image to detect objects using the YOLOv8 model."
)

# Launch the Gradio app
if __name__ == "__main__":
    gr_interface.launch()