Update app.py

app.py

@@ -3,8 +3,7 @@ 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
+from matplotlib.patches import Patch
 
 # Cargar el modelo YOLO (asegúrate de que 'model.pt' esté en el mismo directorio)
 model = YOLO("model.pt")
@@ -103,179 +102,15 @@ def process_image(image):
                 for i in range(3):
                     mask_image[:, :, i][unique_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_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_bgr.copy()
-            print("No se detectaron máscaras en esta imagen.")
-
-    # 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
-    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:
-        handles = []
-        labels = []
-        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"{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_figure
-
-# Crear la interfaz de Gradio
-iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="numpy"),
-    title="Detección de Estenosis",
-    description="Sube una imagen para detectar estenosis."
-)
-
-if __name__ == "__main__":
-    iface.launch()
-import cv2
-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 (asegúrate de que 'model.pt' esté en el mismo directorio)
-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)
-    img_rgb = np.array(image)
-    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
-
-    # 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 (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
-        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()  # 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
-
-            # 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
-            resized_masks = []
-            for mask in masks:
-                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)  # Shape: (num_masks, height, width)
-
-            # 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)
-                # 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 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 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
-
-            # 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))
-
-                # 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]
-                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_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][unique_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_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_bgr.copy()
-            print("No se detectaron máscaras en esta imagen.")
+    # Superponer la imagen de máscaras sobre la imagen original
+    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
     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))
+    fig, ax = plt.subplots(figsize=(10, 10))
     ax.imshow(img_with_masks_rgb)
     ax.axis('off')
 
@@ -285,13 +120,13 @@ def process_image(image):
         labels = []
         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)
+            patch = Patch(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')
+        legend = ax.legend(handles, labels, loc='upper right', bbox_to_anchor=(1.15, 1))
         ax.add_artist(legend)
 
     plt.tight_layout()