Update app.py

app.py

@@ -3,14 +3,12 @@ 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 (PIL usa RGB, OpenCV usa BGR)
+    # Convertir la imagen de PIL a NumPy array y de RGB a BGR
     img_rgb = np.array(image)
     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
 
@@ -63,7 +61,7 @@ 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')  # Cambia 'viridis' por 'plasma' o 'inferno' si lo deseas
+            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)
@@ -106,45 +104,32 @@ def process_image(image):
             img_with_masks = img_bgr.copy()
             print("No se detectaron máscaras en esta imagen.")
 
-    # Convertir la imagen de BGR a RGB para matplotlib
+    # Convertir las imágenes de BGR a RGB para mostrarlas en Gradio
+    img_original_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
     img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
 
-    # 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
+    # Generar la leyenda como texto
     if mask_info_list:
-        handles = []
-        labels = []
+        legend_text = ""
         for mask_info in mask_info_list:
-            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"Stenosis probability: {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
+            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."
 
-    return img_figure
+    return img_original_rgb, img_with_masks_rgb, legend_text
 
-# Crear la interfaz de Gradio
+# Crear la interfaz de Gradio con múltiples salidas
 iface = gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="numpy"),
+    outputs=[
+        gr.Image(type="numpy", label="Imagen Original"),
+        gr.Image(type="numpy", label="Imagen con Máscaras"),
+        gr.HTML(label="Leyenda")
+    ],
     title="Detección de Estenosis",
     description="Sube una imagen para detectar estenosis."
 )