import gradio as gr
import matplotlib.pyplot as plt
from PIL import Image
from ultralyticsplus import YOLO, render_result
import cv2
import numpy as np


model = YOLO('best (1).pt')

def response(image):
  print(image)
  results = model(image)
  for i, r in enumerate(results):
      
    # Plot results image
      im_bgr = r.plot()  
      im_rgb = im_bgr[..., ::-1]  # Convert BGR to RGB
    
    # im_rgb = Image.fromarray(im_rgb)
      
      return im_rgb

def yoloV8_func(image: gr.Image = None,
                image_size: gr.Slider = 640,
                conf_threshold: gr.Slider = 0.4,
                iou_threshold: gr.Slider = 0.50):
    
   
    

    # Perform object detection on the input image using the YOLOv8 model
    results = model.predict(image,
                            conf=conf_threshold,
                            iou=iou_threshold,
                            imgsz=image_size)

    # Print the detected objects' information (class, coordinates, and probability)
    box = results[0].boxes
    print("Object type:", box.cls)
    print("Coordinates:", box.xyxy)
    print("Probability:", box.conf)

    # Render the output image with bounding boxes around detected objects
    render = render_result(model=model, image=image, result=results[0], rect_th = 4, text_th = 4)
    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 = "YOLOv8 Custom Object Detection by Uyen Nguyen"


examples = [['one.jpg', 900, 0.5, 0.8],
            ['two.jpg', 1152, 0.05, 0.05],
            ['three.jpg', 1024, 0.25, 0.25],
            ['four.jpg', 832, 0.3, 0.3]]

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

# Launch the Gradio interface in debug mode with queue enabled
yolo_app.launch(debug=True, share=True)


iface = gr.Interface(fn=response, inputs="image", outputs="image")
iface.launch()