import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from tensorflow.keras.layers import DepthwiseConv2D
from tensorflow.keras.models import Model
from PIL import Image

# Función personalizada para DepthwiseConv2D
def custom_depthwise_conv2d(**kwargs):
    kwargs.pop('groups', None)  # Eliminar el argumento 'groups'
    return DepthwiseConv2D(**kwargs)

custom_objects = {'DepthwiseConv2D': custom_depthwise_conv2d}

# Ruta del modelo
path_to_model = "modelo_jeysshon_iaderm.h5"

# Cargar el modelo con custom_objects
try:
    model = tf.keras.models.load_model(path_to_model, custom_objects=custom_objects)
    print("Modelo cargado exitosamente.")
except Exception as e:
    print(f"Error al cargar el modelo: {e}")
    raise

# Imprimir nombres de las capas del modelo para identificar la última capa convolucional
print("Nombres de las capas en el modelo:")
for layer in model.layers:
    print(layer.name)

# Definir las etiquetas
labels = [
    'Acné / Rosácea', 'Queratosis Actínica / Carcinoma Basocelular', 
    'Dermatitis Atópica', 'Enfermedad Bullosa', 
    'Celulitis Impétigo (Infecciones Bacterianas)', 
    'Eczema', 'Exanthems (Erupciones Cutáneas por Medicamentos)', 'Pérdida de Cabello (Alopecia)', 
    'Herpes HPV', 'Trastornos de Pigmentación', 
    'Lupus',
    'Melanoma (Cáncer de Piel)', 'Hongos en las Uñas', 
    'Hiedra Venenosa', 
    'Psoriasis (Lichen Planus)', 'Sarna Lyme', 
    'Queratosis Seborreica', 'Enfermedad Sistémica', 
    'Tinea Ringworm (Infecciones Fúngicas)', 
    'Urticaria Ronchas', 'Tumores Vasculares', 'Vasculitis', 'Verrugas Molusco'
]

# Función Grad-CAM para generar el mapa de calor de activación
def grad_cam(model, img_array, last_conv_layer_name, pred_index=None):
    grad_model = Model(
        [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
    )

    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model(img_array)
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

    # Calcula los gradientes de la salida de la clase respecto a la salida de la última capa convolucional
    grads = tape.gradient(class_channel, last_conv_layer_output)

    # Promediar sobre cada filtro
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

    # Ponderar cada filtro en el mapa de características de la última capa convolucional
    last_conv_layer_output = last_conv_layer_output[0]
    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)

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

# Generar el mapa de calor de activación usando Grad-CAM
def generate_heatmap(image, prediction, last_conv_layer_name="conv5_block3_out"):  # Ajusta el nombre después de revisar las capas
    # Redimensionar la imagen a (224, 224)
    img_array = tf.image.resize(image, (224, 224))
    img_array = tf.expand_dims(img_array, axis=0)  # Añadir una dimensión para el batch
    
    # Obtén el Grad-CAM para la clase predicha
    heatmap = grad_cam(model, img_array, last_conv_layer_name)
    
    # Redimensionar el mapa de calor al tamaño de la imagen
    heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))

    # Convertir el mapa de calor a un rango de 0 a 255 para aplicar color
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

    # Superponer el mapa de calor en la imagen original
    superimposed_image = cv2.addWeighted(np.array(image), 0.6, heatmap, 0.4, 0)

    # Convertir a formato PIL para mostrar en Gradio
    img_pil = Image.fromarray(superimposed_image)
    return img_pil

def classify_image(image):
    # Redimensionar la imagen a (224, 224) antes de la predicción
    image_resized = tf.image.resize(image, (224, 224))
    image_resized = tf.expand_dims(image_resized, axis=0)  # Añadir una dimensión para el batch
    prediction = model.predict(image_resized).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(len(labels))}
    
    # Generar mapa de calor utilizando Grad-CAM
    heatmap_image = generate_heatmap(np.array(image), prediction)
    
    return confidences, heatmap_image

# Configuración de la interfaz de Gradio
title = "AI-DERM DETECTION"
article = (
    "Se propone un sistema automatizado para el diagnóstico de las 23 enfermedades comunes de la piel..."
)
description = (
    "Utilizamos la interfaz de usuario generada por Gradio para ingresar imágenes a nuestra red neuronal..."
)

# Ejemplos de imágenes
examples = [
    ['./123.jpg'],
    ['./acne-closed-comedo-2.jpg'],
    ['./distal-subungual-onychomycosis-86.jpeg'],
    ['./cherry-angioma-16.jpg'],
    ['./malignant-melanoma-16.jpg'],  
    ['./tinea-primary-lesion-15.jpeg'],
    ['./congenital-nevus-35.jpg'],
    ['./tinea-body-137.jpg'],  
    ['./atopic-13.jpg'], 
    ['./atopic-7.jpg'] 
]

# Lanzar la interfaz de Gradio
gr.Interface(
    fn=classify_image,
    title=title,
    article=article,
    description=description,
    inputs=gr.Image(),
    outputs=[gr.Label(num_top_classes=4), gr.Image(label="Mapa de Calor")],
    examples=examples
).launch()