import gradio as gr
import cv2
import requests
import os

from ultralytics import YOLO

file_urls = [
    "https://www.dropbox.com/scl/fi/2mlc191y9lzbe8nwu45ss/duck76.jpeg?rlkey=qwnr78mtdy7sjg71ldrui6vf0&dl=0",
    "https://www.dropbox.com/scl/fi/2y4cdkwwn3drlh4ob86ic/duck85.jpeg?rlkey=lcl3n0jav7ougsj4tamm1hh93&dl=0",
    "https://www.dropbox.com/scl/fi/8gojxnm8wwhs2isj6k4zv/duck23.jpeg?rlkey=2zioinhr0wfpv22qq963tnv8a&dl=0",
]

def download_file(url, save_name):
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, "wb").write(file.content)

for i, url in enumerate(file_urls):
    if "mp4" in url:
        download_file(url, "video.mp4")
    else:
        download_file(url, f"image_{i}.jpg")

model = YOLO("best.pt")

def show_preds_image(image):
    outputs = model.predict(source=image)
    image_np = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    results = outputs.xyxy[0].cpu().numpy()

    for det in results:
        cv2.rectangle(
            image_np,
            (int(det[0]), int(det[1])),
            (int(det[2]), int(det[3])),
            color=(0, 0, 255),
            thickness=2,
            lineType=cv2.LINE_AA,
        )
    return cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)

inputs_image = gr.inputs.Image(type="file", label="Input Image")
outputs_image = gr.outputs.Image(type="numpy", label="Output Image")
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Duck Image Segmentation",
    examples=file_urls,
    allow_flagging=False,
)

interface_image.launch()