Update app.py

app.py

@@ -3,7 +3,8 @@ import numpy as np
 import matplotlib.pyplot as plt
 from ultralytics import YOLO
 import gradio as gr
-from matplotlib.patches import Patch
+from matplotlib.patches import Rectangle
+from matplotlib.legend import Legend
 
 # Cargar el modelo YOLO
 model = YOLO("model.pt")
@@ -16,25 +17,27 @@ def process_image(image):
     # Realizar inferencia en la imagen
     results = model.predict(source=img, save=False)
 
+    # Crear una carpeta para guardar las imágenes individuales (opcional)
+    # output_folder = "mascaras_individuales"
+    # os.makedirs(output_folder, exist_ok=True)
+
     # Inicializar la lista para almacenar la información de las máscaras
     mask_info_list = []
 
     # Crear una imagen en blanco para las máscaras
     mask_image = np.zeros_like(img, dtype=np.uint8)
 
-    # Transparencia ajustada para la superposición de máscaras
-    alpha = 0.4  # Un valor más bajo aumenta la transparencia y reduce la saturación
-
     # Procesar resultados
     for result in results:
+        # Verificar si se detectaron máscaras
         if result.masks is not None and len(result.masks.data) > 0:
-            # Obtener máscaras, probabilidades y clases
-            masks = result.masks.data.cpu().numpy()
-            confidences = result.boxes.conf.cpu().numpy()
+            # Obtener las máscaras, las probabilidades y las clases
+            masks = result.masks.data.cpu().numpy()  # Forma: (num_masks, altura, ancho)
+            confidences = result.boxes.conf.cpu().numpy()  # Probabilidades
             classes = result.boxes.cls.cpu().numpy().astype(int)
-            names = model.names
+            names = model.names  # Nombres de las clases
 
-            # Normalizar probabilidades al rango [0, 1]
+            # Normalizar las probabilidades al rango [0, 1]
             confidences_norm = (confidences - confidences.min()) / (confidences.max() - confidences.min() + 1e-6)
 
             # Redimensionar las máscaras para que coincidan con el tamaño de la imagen
@@ -42,39 +45,47 @@ def process_image(image):
             for mask in masks:
                 mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)
                 resized_masks.append(mask_resized)
-            resized_masks = np.array(resized_masks)
+            resized_masks = np.array(resized_masks)  # Forma: (num_masks, altura, ancho)
 
             # Aplicar suavizado a las máscaras
             smoothed_masks = []
             for mask in resized_masks:
+                # Convertir la máscara a escala de grises (valores entre 0 y 255)
                 mask_uint8 = (mask * 255).astype(np.uint8)
-                blurred_mask = cv2.GaussianBlur(mask_uint8, (5, 5), 0)
+                # Aplicar desenfoque gaussiano
+                blurred_mask = cv2.GaussianBlur(mask_uint8, (7, 7), 0)
+                # Normalizar y convertir de nuevo a rango [0, 1]
                 mask_smoothed = blurred_mask.astype(np.float32) / 255.0
                 smoothed_masks.append(mask_smoothed)
             smoothed_masks = np.array(smoothed_masks)
 
-            # Ordenar máscaras por probabilidad
+            # Ordenar las máscaras por probabilidad descendente
             sorted_indices = np.argsort(-confidences)
             sorted_masks = smoothed_masks[sorted_indices]
             sorted_confidences = confidences[sorted_indices]
             sorted_confidences_norm = confidences_norm[sorted_indices]
             sorted_classes = classes[sorted_indices]
 
-            # Definir mapa de colores con un ajuste para colores más suaves
+            # Definir el mapa de colores
             colormap = plt.cm.get_cmap('viridis')
 
-            # Procesar cada máscara
-            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
-                update_mask = np.logical_and(mask_bool, mask_image[:, :, 0] == 0)  # Solo actualizamos píxeles que no han sido coloreados
+            # 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)
 
+            # Procesar cada máscara y asignar máscaras de mayor probabilidad a los píxeles
+            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
+                    continue  # Si no hay píxeles nuevos, continuar
+                mask_indices[update_mask] = idx
 
-                # Obtener color y hacerlo más suave
+                # Obtener el color del mapa de colores basado en la probabilidad normalizada
                 color_rgb = colormap(conf_norm)[:3]
-                color_rgb = [int(c * 150) for c in color_rgb]  # Reducir la intensidad del color (150 en lugar de 255)
-                color_bgr = color_rgb[::-1]  # Convertir de RGB a BGR para OpenCV
+                color_rgb = [int(c * 255) for c in color_rgb]  # Convertir a escala [0, 255]
+                color_bgr = color_rgb[::-1]  # Convertir de RGB a BGR
 
                 # Almacenar la información de la máscara
                 mask_info = {
@@ -88,20 +99,34 @@ 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[i]
-
-            # Superponer la imagen de máscaras sobre la imagen original con transparencia
+                    mask_image[:, :, i][update_mask] = color_bgr[i]  # Usar color BGR para OpenCV
+
+                # Generar y guardar la imagen individual de la máscara (opcional)
+                # mask_rgb = np.zeros_like(img, dtype=np.float32)
+                # for i in range(3):
+                #     mask_rgb[:, :, i] = mask * color_bgr[i]  # Usar la máscara suavizada
+                # img_individual = cv2.addWeighted(img.astype(np.float32), 1, mask_rgb, 0.2, 0).astype(np.uint8)
+                # text_str = f"Stenosis probability: {conf:.2f}"
+                # cv2.putText(img_individual, text_str, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color_bgr, 2)
+                # output_path = os.path.join(output_folder, f"imagen_mascara_{idx}.jpg")
+                # cv2.imwrite(output_path, img_individual)
+
+            # Superponer la imagen de máscaras sobre la imagen original
+            alpha = 0.2  # Transparencia ajustada
             img_with_masks = cv2.addWeighted(img.astype(np.float32), 1, mask_image.astype(np.float32), alpha, 0).astype(np.uint8)
 
         else:
+            # Si no hay máscaras, usar la imagen original
             img_with_masks = img.copy()
+            print("No se detectaron máscaras en esta imagen.")
 
-    # Convertir la imagen a RGB para matplotlib
-    img_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
+    # 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)
 
     # Crear una figura para mostrar la imagen y la leyenda
     fig, ax = plt.subplots(figsize=(8, 8))
-    ax.imshow(img_rgb)
+    ax.imshow(img_with_masks_rgb)
     ax.axis('off')
 
     # Crear la leyenda si hay máscaras detectadas
@@ -109,13 +134,18 @@ def process_image(image):
         handles = []
         labels = []
         for mask_info in mask_info_list:
-            color_rgb = np.array(mask_info['color_rgb']) / 255  # Normalizar el color
-            patch = Patch(facecolor=color_rgb)
-            label = f"{mask_info['class']}: {mask_info['confidence']:.2f}"
+            color_rgb = mask_info['color_rgb']
+            color_normalized = np.array(color_rgb) / 255
+            patch = Rectangle((0, 0), 1, 1, facecolor=color_normalized)
+            label = f"Stenosis probability: {mask_info['confidence']:.2f}"
             handles.append(patch)
             labels.append(label)
 
-        ax.legend(handles, labels, loc='upper right')
+        # 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()