import cv2
from ultralytics import YOLO
import gradio as gr

def predict(path: str):
    # Cargar el modelo
    model = YOLO("best4SSL.pt")  
    
    # Leer la imagen
    imagen = cv2.imread(path)
    
    # Reescalar la imagen a 640x640
    imagen_rescaled = cv2.resize(imagen, (640, 640))
    
    # Realizar la predicción sobre la imagen reescalada
    results = model.predict(source=imagen_rescaled)
    
    # Mostrar los resultados
    for r in results:
        return r.plot()  # Asegúrate de que 'r' tiene el método plot()

# Configurar la interfaz Gradio
gr.Interface(fn=predict,
             inputs=gr.components.Image(type="filepath", label="Input"),
             outputs=gr.components.Image(type="numpy", label="Output")).launch()