Update app.py

app.py

@@ -1,21 +1,71 @@
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from ultralytics import YOLO
 import gradio as gr
+import torch
 
 # Cargar un modelo YOLOv8n preentrenado
 model = YOLO('best.pt')
 
+def contar_detecciones(cls_tensor, nombres_clases):
+    conteos = {nombre: torch.sum(cls_tensor == indice).item() for indice, nombre in enumerate(nombres_clases)}
+    return conteos
+
+def clases_detectadas(res):
+    if res and hasattr(res[0], 'xyxy'):
+        cls_tensor = res[0].xyxy[0][:, -1]  # Obtener tensor de clases
+        nombres_clases = model.names
+        conteos = contar_detecciones(cls_tensor, nombres_clases)
+        respuesta = ""
+        for nombre, conteo in conteos.items():
+            respuesta += f"Clase {nombre} : {conteo} detecciones\n"
+        return respuesta
+    else:
+        return "No se encontraron resultados o falta información relevante"
+
 def detect_objects(image: Image.Image):
     # Realizar la inferencia
     results = model.predict(image)
 
-    # Obtener los resultados y el texto de descripción
-    description = ""
-    for r in results:
-        im_array = r.plot()  # plot a BGR numpy array of predictions
-        im = Image.fromarray(im_array[..., ::-1])  # Convertir a imagen RGB
+    # Guardar la imagen con todas las detecciones
+    im_array = results.render()[0]
+    im_all_detections = Image.fromarray(im_array[..., ::-1])
+
+    # Contar las clases detectadas
+    conteo_clases = clases_detectadas(results)
+
+    # Guardar información de detección
+    detections = results.xyxy[0].tolist()
+
+    return im_all_detections, conteo_clases, detections
+
+def update_image(original_image: Image.Image, detections, show_potatoes: bool, show_tongues: bool):
+    # Crear una copia de la imagen original
+    updated_image = original_image.copy()
+    draw = ImageDraw.Draw(updated_image)
+
+    # Definir la fuente para las etiquetas
+    try:
+        font = ImageFont.truetype("arial.ttf", 15)
+    except IOError:
+        font = ImageFont.load_default()
+
+    # Filtrar y dibujar solo las detecciones seleccionadas
+    for det in detections:
+        label = model.names[int(det[5])]
+        if (label == 'papa' and show_potatoes) or (label == 'lengua' and show_tongues):
+            box = det[:4]
+            label_text = f"{label} {det[4]:.2f}"
+            draw.rectangle(box, outline="red", width=2)
+            text_size = draw.textsize(label_text, font=font)
+            draw.rectangle([box[0], box[1] - text_size[1], box[0] + text_size[0], box[1]], fill="red")
+            draw.text((box[0], box[1] - text_size[1]), label_text, fill="white", font=font)
 
-    return im
+    return updated_image
 
 # Crear la interfaz de Gradio
-gr.Interface(fn=detect_objects, inputs="image", outputs="image").launch()
+iface = gr.Interface(
+    fn=detect_objects,
+    update=update_image,
+    inputs=["image", gr.Checkbox(label="Mostrar Papas"), gr.Checkbox(label="Mostrar Lenguas")],
+    outputs=["image", "text", "image"]
+).launch()