app.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
from ultralytics import YOLO
import gradio as gr
from matplotlib.patches import Patch

# 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 una matriz para almacenar la máxima confianza por píxel
    max_confidence_map = np.zeros((img_bgr.shape[0], img_bgr.shape[1]), dtype=np.float32)

    # Definir el mapa de colores (puedes cambiar a 'plasma', 'inferno', etc.)
    colormap = plt.cm.get_cmap('tab20')  # 'tab20' proporciona una variedad de colores distintivos

    # 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]

            # Asignar un color único a cada máscara utilizando el mapa de colores
            for idx_in_order, (mask, conf_norm, conf, cls) in enumerate(
                    zip(sorted_masks, sorted_confidences_norm, sorted_confidences, sorted_classes)):
                
                # Obtener el color del mapa de colores
                color = colormap(idx_in_order % colormap.N)[:3]  # Evitar exceder el número de colores disponibles
                color_rgb = [int(c * 255) for c in color]  # Escalar a [0, 255]
                color_bgr = color_rgb[::-1]  # Convertir de RGB a BGR

                # Umbralizar la máscara para obtener valores binarios
                mask_binary = mask > 0.5

                # Obtener los píxeles donde la máscara actual tiene mayor confianza
                update_pixels = mask_binary & (sorted_confidences[idx_in_order] > max_confidence_map)

                # Actualizar la imagen de máscaras y el mapa de máxima confianza
                for i in range(3):
                    mask_image[:, :, i][update_pixels] = color_bgr[i]
                max_confidence_map[update_pixels] = sorted_confidences[idx_in_order]

                # Almacenar la información de la máscara para la leyenda
                mask_info = {
                    'class': names[cls],
                    'confidence': sorted_confidences[idx_in_order],
                    'color_rgb': color_rgb
                }
                mask_info_list.append(mask_info)

    # 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 y Gradio
    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=(10, 10))
    ax.imshow(img_with_masks_rgb)
    ax.axis('off')

    # Crear la leyenda si hay máscaras detectadas
    if mask_info_list:
        handles = []
        labels = []
        # Para evitar duplicados en la leyenda (si hay múltiples máscaras de la misma clase)
        seen = set()
        for mask_info in mask_info_list:
            label = f"{mask_info['class']} - Confianza: {mask_info['confidence']:.2f}"
            if label not in seen:
                seen.add(label)
                color_rgb_normalized = np.array(mask_info['color_rgb']) / 255  # Normalizar al rango [0, 1]
                patch = Patch(facecolor=color_rgb_normalized)
                handles.append(patch)
                labels.append(label)

        # Añadir la leyenda al gráfico
        ax.legend(handles, labels, loc='upper right', bbox_to_anchor=(1.15, 1))

    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()