import gradio as gr
import numpy as np
import cv2
from keras.models import load_model

import tensorflow as tf
from tensorflow import keras

# Cargar el modelo entrenado
model1 = load_model('./isatron_v3.h5')

# Función para encontrar la última capa convolucional
def find_last_conv_layer(model):
    for layer in reversed(model.layers):
        if 'conv' in layer.name:
            return layer.name
    raise ValueError("No se encontró una capa convolucional en el modelo.")

# Obtener el nombre de la última capa convolucional
last_conv_layer_name = find_last_conv_layer(model1)
print("Última capa convolucional:", last_conv_layer_name)

# Definir tamaño de imagen y etiquetas
img_size1 = 150
labels = ['PNEUMONIA', 'NORMAL']

def load_and_preprocess_image1(img):
    # Convertir imagen de Gradio (PIL Image) a array numpy
    img = np.array(img)
    # Convertir de RGB a escala de grises
    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    # Redimensionar imagen al tamaño requerido
    img = cv2.resize(img, (img_size1, img_size1))
    # Reformatear imagen para entrada del modelo
    img = img.reshape(-1, img_size1, img_size1, 1)
    # Normalizar imagen
    img = img / 255.0
    return img

def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    # Crear un modelo que mapee la imagen de entrada a las activaciones
    # de la última capa convolucional y las predicciones
    grad_model = keras.models.Model(
        [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
    )

    # Calcular el gradiente de la clase predicha con respecto a las activaciones
    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model(img_array)
        if pred_index is None:
            pred_index = np.argmax(preds[0])
        class_channel = preds[:, pred_index]

    # Calcular los gradientes
    grads = tape.gradient(class_channel, last_conv_layer_output)

    # Pooling global de los gradientes
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

    # Multiplicar cada canal por su importancia
    last_conv_layer_output = last_conv_layer_output[0]
    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)

    # Normalizar el mapa de calor entre 0 y 1
    heatmap = tf.maximum(heatmap, 0) / tf.reduce_max(heatmap)
    heatmap = heatmap.numpy()
    return heatmap

def overlay_heatmap(heatmap, img, alpha=0.4):
    # Redimensionar mapa de calor al tamaño de la imagen
    heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
    # Convertir mapa de calor a RGB
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    # Convertir imagen a BGR
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    # Aplicar mapa de calor a la imagen original
    overlayed_img = heatmap * alpha + img
    overlayed_img = np.uint8(overlayed_img)
    # Convertir de nuevo a RGB
    overlayed_img = cv2.cvtColor(overlayed_img, cv2.COLOR_BGR2RGB)
    return overlayed_img

def image_classifier1(img):
    # Mantener la imagen original para superponer
    orig_img = np.array(img)
    # Preprocesar la imagen
    img_array = load_and_preprocess_image1(img)
    # Realizar predicción usando model1
    preds = model1.predict(img_array)
    prediction = preds[0][0]  # Suponiendo que el modelo devuelve una probabilidad
    # Determinar el índice de la clase predicha
    pred_index = int(prediction > 0.5)
    # Generar mapa de calor
    heatmap = make_gradcam_heatmap(img_array, model1, last_conv_layer_name, pred_index=pred_index)
    # Superponer mapa de calor en la imagen original
    overlayed_img = overlay_heatmap(heatmap, orig_img)
    # Retornar la imagen superpuesta y los porcentajes de predicción
    prediction_percentage = {'PNEUMONIA': float(prediction), 'NORMAL': float(1 - prediction)}
    return overlayed_img, prediction_percentage

# Crear interfaz Gradio
demo_model1 = gr.Interface(
    fn=image_classifier1,
    inputs=gr.Image(type="pil", label="Subir imagen de rayos X"),
    outputs=[gr.Image(type="numpy", label="Imagen con Mapa de Calor"), gr.Label(label="Predicción")],
    title="<h1 style='text-align: center;'>IsaTron V2: Herramienta de Apoyo al Diagnóstico de Neumonía</h1>",
    description="""
    <div style='text-align: justify;'>
    IsaTron es una herramienta de inteligencia artificial desarrollada con redes neuronales convolucionales (CNN) para apoyar el diagnóstico de neumonía a partir de imágenes de rayos X. La IA analiza las imágenes y produce un mapa de calor que resalta las áreas de mayor relevancia para la predicción del modelo, lo cual ayuda a los profesionales de la salud a visualizar mejor las zonas afectadas.
    </div>
    <div style='text-align: justify;'>
    <strong>Advertencia:</strong> IsaTron está diseñado como una herramienta de apoyo y no reemplaza una evaluación médica profesional. Es crucial que los resultados generados por esta herramienta sean interpretados por personal de salud calificado.
    </div>
    """,
    article="""
    <div style='text-align: justify;'>
    Este proyecto sigue en desarrollo y utiliza tecnologías avanzadas de inteligencia artificial, como redes neuronales convolucionales y Grad-CAM, para mejorar la interpretabilidad de los resultados. IsaTron ha sido entrenado con imágenes médicas y es capaz de predecir neumonía con un alto grado de confianza. Sin embargo, los resultados obtenidos deben ser confirmados por un médico especialista para realizar un diagnóstico clínico adecuado.
    </div>

    <div style='text-align: center;'>Jeysshon, 2024.</div>
    """
)

# Ejecutar la interfaz
if __name__ == "__main__":
    demo_model1.launch(share=True)