from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator
import torch
import gradio as gr
import cv2
model = YOLO('best.pt')
def yolo_pred(image):
    result = model.predict(image)[0]
    annotator = Annotator(result.orig_img)
    color_list = [(107, 31, 45), (32, 102, 50), (32, 45, 102)]
    for label in result.boxes.data.detach().numpy():
        annotator.box_label(
            label[0:4],
            str(result.names[label[-1].item()]) + " " + str(round(label[-2], 2)),
            color_list[int(label[-1].item())]
        )
        print(round(label[-2], 2))
    return annotator.im
gr.Interface(fn=yolo_pred, 
             inputs="image", 
             outputs="image",
             examples=[
                 [cv2.imread("example1.jpg")],
                 [cv2.imread("example2.jpg")],
                 [cv2.imread("example3.jpg")],
             ],
             title="Fine-Tuned YOLOv8",
             description="""YOLOv8 object detection model trained on the Tsinghua-Daimler Cyclist Benchmark (TDCB). 
                             Since the setting of their image collection seems to be an early morning in China,
                             please sample similar images for the best results. I recommend using images from TDCB's
                             Kaggle clone here: https://www.kaggle.com/datasets/semiemptyglass/cyclist-dataset."""
            ).launch()