import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
def draw_boxes(image, boxes):
    for box in boxes:        
        x1, y1, x2, y2, name, prob = box
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, f"{name} {prob:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1.9, (255,55,0), 2)  
    return image
def detect_objects_on_image(buf):    
    model = YOLO("files/yolov8n.pt")      
    results = model.predict(buf)
    result = results[0]
    output = []
    for box in result.boxes:
        x1, y1, x2, y2 = [
            round(x) for x in box.xyxy[0].tolist()
        ]
        class_id = box.cls[0].item()
        prob = round(box.conf[0].item(), 2)
        output.append([
            x1, y1, x2, y2, result.names[class_id], prob
        ])
    return output

path  = [['files/a.png'],['files/image_0.png'],['files/huggingface.png']]
inputs = [
gr.inputs.Image(type="filepath", label="Input Image"),
gr.inputs.Radio(label="YOLO Methods", 
                choices=[                         
                        "v8"                        
                    ]),  
   
]
outputs = [   
gr.outputs.Image(type="numpy", label="Output YOLO Image")     
]
def process_image(input_image, radio_choice,slider_val):    
    img = cv2.imread(input_image)  
    img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB) 
    boxes=detect_objects_on_image(img)
    img_with_boxes = draw_boxes(img, boxes)                  
    return img_with_boxes
def on_change(radio_choice):
    outputs[0].update(process_image(
        inputs[0].value,          
        radio_choice)
    )
yolo = gr.Interface(
    fn=process_image,
    inputs=inputs,
    outputs=outputs, 
    on_change=on_change,
    examples=path,
    title="YOLO: Computer Vision and Deep Learning by Farshid PirahanSiah",    
    )