Update app.py

app.py

@@ -3,12 +3,14 @@ import numpy as np
 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
 model = YOLO("model.pt")
 
 def process_image(image):
-    # Convertir la imagen de PIL a NumPy array y de RGB a BGR
+    # Convertir la imagen de PIL a NumPy array y de RGB a BGR (PIL usa RGB, OpenCV usa BGR)
     img_rgb = np.array(image)
     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
 
@@ -21,6 +23,9 @@ def process_image(image):
     # Crear una imagen en blanco para las máscaras (en formato BGR)
     mask_image = np.zeros_like(img_bgr, dtype=np.uint8)
 
+    # Inicializar la máscara acumulativa
+    cumulative_mask = np.zeros((img_bgr.shape[0], img_bgr.shape[1]), dtype=bool)
+
     # Procesar resultados
     for result in results:
         # Verificar si se detectaron máscaras
@@ -61,20 +66,23 @@ def process_image(image):
             sorted_classes = classes[sorted_indices]
 
             # Definir el mapa de colores (puedes cambiar a 'plasma', 'inferno', etc.)
-            colormap = plt.cm.get_cmap('viridis')
-
-            # Crear una matriz para rastrear qué máscara se asigna a cada píxel
-            mask_indices = np.full((img_bgr.shape[0], img_bgr.shape[1]), -1, dtype=int)
+            colormap = plt.cm.get_cmap('viridis')  # Cambia 'viridis' por 'plasma' o 'inferno' si lo deseas
 
-            # Procesar cada máscara y asignar máscaras de mayor probabilidad a los píxeles
+            # Procesar cada máscara y asignar máscaras de mayor confianza primero
             for idx_in_order, (idx, mask, conf_norm, conf, cls) in enumerate(
                     zip(sorted_indices, sorted_masks, sorted_confidences_norm, sorted_confidences, sorted_classes)):
-                mask_bool = mask > 0.5  # Umbral para convertir a binario
-                # Actualizar píxeles donde aún no se ha asignado una máscara
-                update_mask = np.logical_and(mask_bool, mask_indices == -1)
-                if not np.any(update_mask):
-                    continue  # Si no hay píxeles nuevos, continuar
-                mask_indices[update_mask] = idx
+                # Umbralizar la máscara para obtener valores binarios
+                mask_bool = mask > 0.5
+
+                # Restar la máscara acumulativa de la máscara actual para obtener la parte única
+                unique_mask = np.logical_and(mask_bool, np.logical_not(cumulative_mask))
+
+                # Actualizar la máscara acumulativa
+                cumulative_mask = np.logical_or(cumulative_mask, unique_mask)
+
+                # Si no hay píxeles únicos, continuar
+                if not np.any(unique_mask):
+                    continue
 
                 # Obtener el color del mapa de colores basado en la probabilidad normalizada
                 color_rgb = colormap(conf_norm)[:3]  # Color en formato RGB [0, 1]
@@ -93,7 +101,7 @@ def process_image(image):
 
                 # Asignar colores a los píxeles correspondientes en la imagen de máscaras
                 for i in range(3):
-                    mask_image[:, :, i][update_mask] = color_bgr_255[i]
+                    mask_image[:, :, i][unique_mask] = color_bgr_255[i]
 
             # Superponer la imagen de máscaras sobre la imagen original
             alpha = 0.2  # Transparencia ajustada
@@ -104,32 +112,45 @@ def process_image(image):
             img_with_masks = img_bgr.copy()
             print("No se detectaron máscaras en esta imagen.")
 
-    # Convertir las imágenes de BGR a RGB para mostrarlas en Gradio
-    img_original_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    # Convertir la imagen de BGR a RGB para matplotlib
     img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
 
-    # Generar la leyenda como texto
+    # Crear una figura para mostrar la imagen y la leyenda
+    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:
-        legend_text = ""
+        handles = []
+        labels = []
         for mask_info in mask_info_list:
-            color_rgb = mask_info['color_rgb']
-            color_hex = '#%02x%02x%02x' % tuple(color_rgb)
-            label = f"<span style='color:{color_hex};'>■</span> {mask_info['class']}: {mask_info['confidence']:.2f}<br>"
-            legend_text += label
-    else:
-        legend_text = "No se detectaron máscaras en esta imagen."
+            color_rgb_normalized = np.array(mask_info['color_rgb']) / 255  # Normalizar al rango [0, 1]
+            patch = Rectangle((0, 0), 1, 1, facecolor=color_rgb_normalized)
+            label = f"{mask_info['class']} - Confianza: {mask_info['confidence']:.2f}"
+            handles.append(patch)
+            labels.append(label)
+
+        # Añadir la leyenda al gráfico
+        legend = Legend(ax, handles, labels, loc='upper right')
+        ax.add_artist(legend)
+
+    plt.tight_layout()
+
+    # Convertir la figura a una imagen NumPy
+    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_original_rgb, img_with_masks_rgb, legend_text
+    return img_figure
 
-# Crear la interfaz de Gradio con múltiples salidas
+# Crear la interfaz de Gradio
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Image(type="numpy", label="Imagen Original"),
-        gr.Image(type="numpy", label="Imagen con Máscaras"),
-        gr.HTML(label="Leyenda")
-    ],
+    outputs=gr.Image(type="numpy"),
     title="Detección de Estenosis",
     description="Sube una imagen para detectar estenosis."
 )