import gradio as gr
from transformers import AutoImageProcessor, AutoModelForObjectDetection, pipeline

# Carga el procesador de imágenes y el modelo
image_processor = AutoImageProcessor.from_pretrained("seayala/practica_2")
model = AutoModelForObjectDetection.from_pretrained("seayala/practica_2")

# Crea el pipeline de detección de objetos
detector = pipeline("object-detection", model=model, image_processor=image_processor)

# Función para procesar la imagen y generar anotaciones
def predict(image):
    results = detector(image)

    # Extrae cajas en formato xmin, ymin, xmax, ymax
    boxes = []
    for obj in results:
        box = obj["box"]
        label = f'{obj["label"]} ({obj["score"]:.2f})'
        # Convierte el formato si es necesario
        if "x" in box and "y" in box and "width" in box and "height" in box:
            xmin = box["x"]
            ymin = box["y"]
            xmax = xmin + box["width"]
            ymax = ymin + box["height"]
        else:
            xmin = box.get("xmin", 0)
            ymin = box.get("ymin", 0)
            xmax = box.get("xmax", 0)
            ymax = box.get("ymax", 0)

        boxes.append({"label": label, "box": [xmin, ymin, xmax, ymax]})

    return image, boxes

# Interfaz Gradio
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Sube una imagen"),
    outputs=gr.AnnotatedImage(label="Resultados de detección"),
    title="Detector de objetos",
    description="Sube una imagen para detectar objetos con tu modelo personalizado."
)

iface.launch()