Update app.py

app.py

@@ -1,27 +1,45 @@
 import gradio as gr
 from transformers import AutoImageProcessor, AutoModelForObjectDetection, pipeline
-import torch
 
 # Carga el procesador de imágenes y el modelo
-image_processor = AutoImageProcessor.from_pretrained("seayala/practica_2")  # Reemplaza con la ruta de tu modelo
-model = AutoModelForObjectDetection.from_pretrained("seayala/practica_2")  # Reemplaza con la ruta de tu 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)
 
-# Crea la función de predicción
+# Función para procesar la imagen y generar anotaciones
 def predict(image):
-  results = detector(image)
-  return results
+    results = detector(image)
 
-# Crea la interfaz de usuario
+    # 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"),
-      outputs=gr.Label(num_top_classes=5),
-      title="Detector de objetos",
-      description="Sube una imagen para detectar objetos.",
-  )
-
-# Inicia la interfaz de usuario
-iface.launch()
+    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()