Update app.py

app.py

@@ -4,7 +4,6 @@ import matplotlib.pyplot as plt
 from ultralytics import YOLO
 import gradio as gr
 from matplotlib.patches import Rectangle
-from matplotlib.legend import Legend
 
 # Cargar el modelo YOLO (asegúrate de que 'model.pt' esté en el mismo directorio)
 model = YOLO("model.pt")
@@ -20,7 +19,10 @@ def process_image(image):
     # Inicializar la lista para almacenar la información de las máscaras
     mask_info_list = []
 
-    # Iterar sobre los resultados
+    # Crear una imagen en blanco para las máscaras
+    mask_image = np.zeros_like(img, dtype=np.uint8)
+
+    # Procesar resultados
     for result in results:
         # Verificar si se detectaron máscaras
         if result.masks is not None and len(result.masks.data) > 0:
@@ -62,9 +64,6 @@ def process_image(image):
             # Definir el mapa de colores
             colormap = plt.cm.get_cmap('viridis')
 
-            # Crear una imagen en blanco para las máscaras
-            mask_image = np.zeros_like(img, dtype=np.uint8)
-
             # Crear una matriz para rastrear qué máscara se asigna a cada píxel
             mask_indices = np.full((img.shape[0], img.shape[1]), -1, dtype=int)
 
@@ -105,10 +104,39 @@ def process_image(image):
             img_with_masks = img.copy()
             print("No se detectaron máscaras en esta imagen.")
 
-    # Convertir la imagen de BGR a RGB para mostrarla con Gradio
+    # Convertir las imágenes de BGR a RGB para matplotlib
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
 
-    return img_with_masks_rgb
+    # Crear una figura de matplotlib
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.imshow(img_with_masks_rgb)
+    ax.axis('off')
+
+    # Crear la leyenda si hay máscaras detectadas
+    if mask_info_list:
+        from matplotlib.patches import Patch
+        handles = []
+        labels = []
+        for mask_info in mask_info_list:
+            color_rgb = mask_info['color_rgb']
+            color_normalized = np.array(color_rgb) / 255
+            patch = Patch(facecolor=color_normalized)
+            label = f"{mask_info['class']}: {mask_info['confidence']:.2f}"
+            handles.append(patch)
+            labels.append(label)
+
+        # Añadir la leyenda al gráfico
+        ax.legend(handles, labels, loc='upper right')
+
+    # Convertir la figura de matplotlib a una imagen que pueda ser mostrada en Gradio
+    fig.canvas.draw()
+    img_figure = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    img_figure = img_figure.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+
+    plt.close(fig)  # Cerrar la figura para liberar memoria
+
+    return img_figure
 
 # Crear la interfaz de Gradio
 iface = gr.Interface(