Update app.py

app.py

@@ -22,8 +22,11 @@ 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)
+    # Inicializar una matriz para almacenar la máxima confianza por píxel
+    max_confidence_map = np.zeros((img_bgr.shape[0], img_bgr.shape[1]), dtype=np.float32)
+
+    # Definir el mapa de colores (puedes cambiar a 'plasma', 'inferno', etc.)
+    colormap = plt.cm.get_cmap('tab20')  # 'tab20' proporciona una variedad de colores distintivos
 
     # Procesar resultados
     for result in results:
@@ -64,49 +67,39 @@ def process_image(image):
             sorted_confidences_norm = confidences_norm[sorted_indices]
             sorted_classes = classes[sorted_indices]
 
-            # Definir el mapa de colores (puedes cambiar a 'plasma', 'inferno', etc.)
-            colormap = plt.cm.get_cmap('plasma')  # Cambia 'viridis' por 'plasma' o 'inferno' si lo deseas
+            # Asignar un color único a cada máscara utilizando el mapa de colores
+            for idx_in_order, (mask, conf_norm, conf, cls) in enumerate(
+                    zip(sorted_masks, sorted_confidences_norm, sorted_confidences, sorted_classes)):
+                
+                # Obtener el color del mapa de colores
+                color = colormap(idx_in_order % colormap.N)[:3]  # Evitar exceder el número de colores disponibles
+                color_rgb = [int(c * 255) for c in color]  # Escalar a [0, 255]
+                color_bgr = color_rgb[::-1]  # Convertir de RGB a BGR
 
-            # 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)):
                 # 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))
+                mask_binary = mask > 0.5
 
-                # Actualizar la máscara acumulativa
-                cumulative_mask = np.logical_or(cumulative_mask, unique_mask)
+                # Obtener los píxeles donde la máscara actual tiene mayor confianza
+                update_pixels = mask_binary & (sorted_confidences[idx_in_order] > max_confidence_map)
 
-                # Si no hay píxeles únicos, continuar
-                if not np.any(unique_mask):
-                    continue
+                # Actualizar la imagen de máscaras y el mapa de máxima confianza
+                for i in range(3):
+                    mask_image[:, :, i][update_pixels] = color_bgr[i]
+                max_confidence_map[update_pixels] = sorted_confidences[idx_in_order]
 
-                # Obtener el color del mapa de colores basado en la probabilidad normalizada
-                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
+                # Almacenar la información de la máscara para la leyenda
                 mask_info = {
-                    'mask_index': idx,
                     'class': names[cls],
-                    'confidence': conf,
-                    'color_rgb': color_rgb_255,
-                    'color_bgr': color_bgr_255
+                    'confidence': sorted_confidences[idx_in_order],
+                    'color_rgb': color_rgb
                 }
                 mask_info_list.append(mask_info)
 
-                # Asignar colores a los píxeles correspondientes en la imagen de máscaras
-                # Utilizamos broadcasting para asignar el color a los píxeles únicos
-                mask_image[unique_mask] = color_bgr_255
-
     # Superponer la imagen de máscaras sobre la imagen original
-    alpha = 0.3  # Transparencia ajustada (puedes ajustar este valor según tus preferencias)
+    alpha = 0.4  # Transparencia ajustada (puedes ajustar este valor según tus preferencias)
     img_with_masks = cv2.addWeighted(img_bgr.astype(np.float32), 1, mask_image.astype(np.float32), alpha, 0).astype(np.uint8)
 
-    # Convertir la imagen de BGR a RGB para matplotlib
+    # Convertir la imagen de BGR a RGB para matplotlib y Gradio
     img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
 
     # Crear una figura para mostrar la imagen y la leyenda
@@ -118,16 +111,19 @@ def process_image(image):
     if mask_info_list:
         handles = []
         labels = []
+        # Para evitar duplicados en la leyenda (si hay múltiples máscaras de la misma clase)
+        seen = set()
         for mask_info in mask_info_list:
-            color_rgb_normalized = np.array(mask_info['color_rgb']) / 255  # Normalizar al rango [0, 1]
-            patch = Patch(facecolor=color_rgb_normalized)
             label = f"{mask_info['class']} - Confianza: {mask_info['confidence']:.2f}"
-            handles.append(patch)
-            labels.append(label)
+            if label not in seen:
+                seen.add(label)
+                color_rgb_normalized = np.array(mask_info['color_rgb']) / 255  # Normalizar al rango [0, 1]
+                patch = Patch(facecolor=color_rgb_normalized)
+                handles.append(patch)
+                labels.append(label)
 
         # Añadir la leyenda al gráfico
-        legend = ax.legend(handles, labels, loc='upper right', bbox_to_anchor=(1.15, 1))
-        ax.add_artist(legend)
+        ax.legend(handles, labels, loc='upper right', bbox_to_anchor=(1.15, 1))
 
     plt.tight_layout()