Update app.py

app.py

@@ -10,29 +10,25 @@ from matplotlib.legend import Legend
 model = YOLO("model.pt")
 
 def process_image(image):
-    # Convertir la imagen de PIL a NumPy array y de RGB a BGR
-    img = np.array(image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    # 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)
 
-    # 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)
+    # Realizar inferencia en la imagen BGR
+    results = model.predict(source=img_bgr, save=False)
 
     # 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)
+    # Crear una imagen en blanco para las máscaras (en formato BGR)
+    mask_image = np.zeros_like(img_bgr, 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:
             # Obtener las máscaras, las probabilidades y las clases
-            masks = result.masks.data.cpu().numpy()  # Forma: (num_masks, altura, ancho)
+            masks = result.masks.data.cpu().numpy()  # Shape: (num_masks, height, width)
             confidences = result.boxes.conf.cpu().numpy()  # Probabilidades
             classes = result.boxes.cls.cpu().numpy().astype(int)
             names = model.names  # Nombres de las clases
@@ -43,9 +39,9 @@ def process_image(image):
             # Redimensionar las máscaras para que coincidan con el tamaño de la imagen
             resized_masks = []
             for mask in masks:
-                mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)
+                mask_resized = cv2.resize(mask, (img_bgr.shape[1], img_bgr.shape[0]), interpolation=cv2.INTER_LINEAR)
                 resized_masks.append(mask_resized)
-            resized_masks = np.array(resized_masks)  # Forma: (num_masks, altura, ancho)
+            resized_masks = np.array(resized_masks)  # Shape: (num_masks, height, width)
 
             # Aplicar suavizado a las máscaras
             smoothed_masks = []
@@ -66,11 +62,11 @@ def process_image(image):
             sorted_confidences_norm = confidences_norm[sorted_indices]
             sorted_classes = classes[sorted_indices]
 
-            # Definir el mapa de colores
-            colormap = plt.cm.get_cmap('viridis')
+            # 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
 
             # 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)
+            mask_indices = np.full((img_bgr.shape[0], img_bgr.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(
@@ -83,45 +79,34 @@ def process_image(image):
                 mask_indices[update_mask] = idx
 
                 # Obtener el color del mapa de colores basado en la probabilidad normalizada
-                color_rgb = colormap(conf_norm)[:3]
-                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
+                color_rgb = colormap(conf_norm)[:3]  # Color en formato RGB [0, 1]
+                color_rgb_255 = [int(c * 255) for c in color_rgb]  # Escalar a [0, 255]
+                color_bgr_255 = color_rgb_255[::-1]  # Convertir de RGB a BGR
 
                 # Almacenar la información de la máscara
                 mask_info = {
                     'mask_index': idx,
                     'class': names[cls],
                     'confidence': conf,
-                    'color_rgb': color_rgb,
-                    'color_bgr': color_bgr
+                    'color_rgb': color_rgb_255,
+                    'color_bgr': color_bgr_255
                 }
                 mask_info_list.append(mask_info)
 
                 # 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]  # 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)
+                    mask_image[:, :, i][update_mask] = color_bgr_255[i]
 
             # 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)
+            img_with_masks = cv2.addWeighted(img_bgr.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()
+            img_with_masks = img_bgr.copy()
             print("No se detectaron máscaras en esta imagen.")
 
-    # Convertir las imágenes de BGR a RGB para matplotlib
-    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    # Convertir la imagen de BGR a RGB para matplotlib
     img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
 
     # Crear una figura para mostrar la imagen y la leyenda
@@ -134,9 +119,8 @@ def process_image(image):
         handles = []
         labels = []
         for mask_info in mask_info_list:
-            color_rgb = mask_info['color_rgb']
-            color_normalized = np.array(color_rgb) / 255
-            patch = Rectangle((0, 0), 1, 1, facecolor=color_normalized)
+            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)