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.") # 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()