#%%
import gradio as gr
from ultralytics import YOLO
import cv2
from huggingface_hub import hf_hub_download

REPO_ID = "RandomCatLover/bird_detector"
FILENAME = "best.pt"

hf_hub_download(repo_id=REPO_ID, filename=FILENAME, local_dir=".")
#%%
yolo = YOLO("best.pt")

def predict(img):
    results = yolo(img)
    for result in results:
        try:
            bbox = result.boxes.xyxy[0]
        except:
            continue
        bbox = [int(i) for i in bbox]
        img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
    return img
#%%
demo = gr.Interface(predict, 
             "image", 
             "image", 
             examples=[
                 f"imgs/{i}.jpg" for i in range(1,7)
                 ],
             cache_examples=True,
             allow_flagging='never')

demo.launch(debug=True)