Create app.py

app.py

@@ -0,0 +1,123 @@
+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
+    img = np.array(image)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+    # Realizar inferencia en la imagen
+    results = model.predict(source=img, save=False)
+
+    # Inicializar la lista para almacenar la información de las máscaras
+    mask_info_list = []
+
+    # Iterar sobre los 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)
+            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.shape[1], img.shape[0]), interpolation=cv2.INTER_LINEAR)
+                resized_masks.append(mask_resized)
+            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)
+                # 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
+            colormap = plt.cm.get_cmap('viridis')
+
+            # Crear una imagen en blanco para las máscaras
+            mask_image = np.zeros_like(img, dtype=np.uint8)
+
+            # 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  # Si no hay píxeles nuevos, continuar
+                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
+
+                # 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
+                }
+                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
+
+            # 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 de BGR a RGB para mostrarla con Gradio
+    img_with_masks_rgb = cv2.cvtColor(img_with_masks, cv2.COLOR_BGR2RGB)
+
+    return img_with_masks_rgb
+
+# 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()