File size: 5,577 Bytes
06f39c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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()
|