import numpy as np
import pandas as pd
import cv2
import gradio as gr
import torch
from ultralyticsplus import YOLO, render_result
import matplotlib.pyplot as plt

# Creating a function to perform predictions
def prediction(image: gr.Image = None, 
                image_size: gr.Slider = 640,
                conf_threshold: gr.Slider = 0.4,
                iou_threshold: gr.Slider = 0.50):

    model = YOLO("best.pt")

    results = model.predict(
        source=image,
        conf=conf_threshold,
        iou=iou_threshold,
        imgsz=image_size
    )
    
    image = cv2.imread(image)
        
    for result in results[0].obb:
        point_1_x = float(result.xyxyxyxy[0][0][0])
        point_1_y = float(result.xyxyxyxy[0][0][1])
        point_2_x = float(result.xyxyxyxy[0][1][0])
        point_2_y = float(result.xyxyxyxy[0][1][1])
        point_3_x = float(result.xyxyxyxy[0][2][0])
        point_3_y = float(result.xyxyxyxy[0][2][1])
        point_4_x = float(result.xyxyxyxy[0][3][0])
        point_4_y = float(result.xyxyxyxy[0][3][1])
        
        cls = int(result.cls)
        
        if cls == 1:
            color = (0, 255, 0)
        else:
            color = (255, 0, 0)
        
        conf = float(result.conf)
        
        text = f"{cls} : {np.round(conf) * 100}%" 
        
        points = np.array([[point_1_x, point_1_y], 
                           [point_2_x, point_2_y], 
                           [point_3_x, point_3_y], 
                           [point_4_x, point_4_y]], np.int32)
         
        points = points.reshape((-1, 1, 2))
        
        cv2.polylines(image, [points], isClosed = True, color = color, thickness = 2)   

    return image

inputs = [
    gr.Image(type="filepath", label="Select an image"),
    gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold")
]

outputs = gr.Image(type = "filepath", label="Output Image")  

yolo_app = gr.Interface(
    fn = prediction,
    inputs = inputs,
    outputs = outputs,
    title = "VPS Model"
)

yolo_app.launch(debug = True, share = True)