import cv2
import numpy as np
import gradio as gr

from ultralytics import YOLO


def predict(path:str, threshold: float = 0.6):
    model = YOLO("best.pt")
    imagen = cv2.imread(path)
    results = model.predict(source=path)

    for r in results:
        # Mantener solo las cajas con una probabilidad mayor al umbral
        boxes = [box for box in r.boxes if box.conf > threshold]
        r.boxes = boxes  # Actualizar las cajas filtradas
        return r.plot()


gr.Interface(fn=predict,
             inputs=gr.components.Image(type="filepath", label="Input"),
             outputs=gr.components.Image(type="numpy", label="Output")).launch(debug=False)