coro_yolo / app.py
fecia's picture
Update app.py
152f86b verified
raw
history blame
13.9 kB
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()
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()