import cv2
import gradio as gr
from ultralytics import YOLO

def inference(path:str, threshold:float=0.6):
    print("trying inference with path", path)
    if path is None:
        return None,0
    model = YOLO('yolo8n_small.pt') # Caution new API since Ultralytics 8.0.43
    outputs = model.predict(source=path, show=False, conf=threshold) # new API with Ultralytics 8.0.43. Accepts 'show', 'classes', 'stream', 'conf' (default is 0.25)
    image = cv2.imread(path)
    counter = 0
    for output in outputs:  # mono item batch
        conf = output.boxes.conf  # the tensor of detection confidences
        xyxy = output.boxes.xyxy
        cls = output.boxes.cls # 0 is 'person' and 5 is 'bus' 16 is dog
        nb=cls.size(dim=0)
        for i in range(nb):
            box = xyxy[i]
            if conf[i]<threshold:
                break
            cv2.rectangle(
                image,
                (int(box[0]), int(box[1])),
                (int(box[2]), int(box[3])),
                color=(0, 0, 255),
                thickness=2,
            )

            # font                   = cv2.FONT_HERSHEY_SIMPLEX
            # fontScale              = 1
            # fontColor              = (255,255,255)
            # thickness              = 1
            # lineType               = 2
            
            # cv2.putText(image,
            # str(cls), 
            # (int(box[0]), int(box[2])), 
            # font, 
            # fontScale,
            # fontColor,
            # thickness,
            # lineType)
            
            counter+=1
    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB), counter

gr.Interface(
    fn = inference,
    inputs = [ gr.components.Image(type="filepath", label="Input"), gr.Slider(minimum=0.0, maximum=0.95, step=0.05, value=0.4, label="Confidence threshold") ],
    outputs = [ gr.components.Image(type="numpy", label="Output"), gr.Label(label="Number of legos detected for given confidence threshold") ],
    title="YOLOv8n LEGO Detection Demo",
    description="LEGO Detection DEMO dataset_small",
    examples=[ ['sample1.jpg'],['sample2.jpg'], ['hard.jpg']],
    allow_flagging="never"
).launch(debug=True, enable_queue=True)