import gradio as gr # for creating interactive UIs in Hugging Face
import torch
from ultralyticsplus import YOLO, render_result # for Hugging Face integration, render_result -> takes the model, image and generates results
 
# Creating a function to perform predictions

def yolov8_predict_func(image : gr.Image = None, # uploading an image as input
                        image_size : gr.Slider = 640, # setting the image size
                        conf_threshold : gr.Slider = 0.4, # setting the confidence threshold
                        iou_threshold : gr.Slider = 0.50): # setting IOU threshold for object detection
        
    # Loading the YOLOv8 model

    model_path = "best.pt"

    model = YOLO(model_path)
    
    # print(model)
    
    # Performing the detection on the YOLO Image
    
    results = model.predict(
        image,
        conf = conf_threshold,
        iou = iou_threshold,
        imgsz = image_size
    )
    
    # Displaying the detected object's information
    
    box = results[0].boxes
    print(f"Object type : {box.cls}")
    print(f"Coordinates : {box.xyxy}")
    print(f"Probability : {box.conf}")
    
    # Rendering the output image with bounding boxes around detected objects
    
    render = render_result(model = model, image = image, result = results[0])

    return render

inputs = [
    gr.Image(type = "filepath", label = "Input 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")

title = "VPS"

examples = [
    ["ps2.0/testing/indoor-parking lot/007.jpg"]
]

yolo_app = gr.Interface(
    fn = yolov8_predict_func,
    inputs = inputs,
    outputs = outputs,
    title = title,
    examples = examples,
    cache_examples = True
)

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