import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2
import datetime  # Importe o módulo datetime

# Defina a camada personalizada FixedDropout
class FixedDropout(tf.keras.layers.Dropout):
    def _get_noise_shape(self, inputs):
        if self.noise_shape is None:
            return self.noise_shape
        symbolic_shape = tf.shape(inputs)
        noise_shape = [symbolic_shape[axis] if shape is None else shape
                       for axis, shape in enumerate(self.noise_shape)]
        return tuple(noise_shape)

# Registre a camada personalizada FixedDropout
tf.keras.utils.get_custom_objects()['FixedDropout'] = FixedDropout

# Carregue seu modelo TensorFlow treinado
with tf.keras.utils.custom_object_scope({'FixedDropout': FixedDropout}):
    model = tf.keras.models.load_model('modelo_treinado.h5')

# Defina uma função para fazer previsões
def classify_image(input_image):
    # Log da forma da entrada
    print(f"Forma da entrada: {input_image.shape}")

    # Redimensione a imagem para as dimensões corretas (192x256)
    input_image = tf.image.resize(input_image, (192, 256))  # Redimensione para as dimensões esperadas
    input_image = (input_image / 255.0)  # Normalize para [0, 1]
    input_image = np.expand_dims(input_image, axis=0)  # Adicione a dimensão de lote

    # Log da forma da entrada após o redimensionamento
    print(f"Forma da entrada após o redimensionamento: {input_image.shape}")

    # Obtenha o tempo atual
    current_time = datetime.datetime.now()

    # Faça a previsão usando o modelo
    prediction = model.predict(input_image)

    # Assumindo que o modelo retorna probabilidades para duas classes, você pode retornar a classe com a maior probabilidade
    class_index = np.argmax(prediction)
    class_labels = ["Normal", "Cataract"]  # Substitua pelas suas etiquetas de classe reais
    predicted_class = class_labels[class_index]

    # Crie uma imagem composta com o rótulo de previsão
    output_image = (input_image[0] * 255).astype('uint8')

    # Set output image dimensions to match input image
    output_image = np.ones((input_image.shape[1], input_image.shape[2], 3), dtype=np.uint8) * 255

    # Add space for the prediction label at the bottom of the image
    label_background = np.ones((50, output_image.shape[1], 3), dtype=np.uint8) * 255

    # Put the text label and box inside the white background
    output_image[-50:] = label_background

    # Write the prediction label on the image
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.4  # Tamanho da fonte reduzido
    cv2.putText(output_image, f"Analysis Time: {current_time.strftime('%Y-%m-%d %H:%M:%S')}", (10, output_image.shape[0] - 30), font, font_scale, (0, 0, 0), 1)
    cv2.putText(output_image, f"Predicted Class: {predicted_class}", (10, output_image.shape[0] - 10), font, font_scale, (0, 0, 0), 1)  # Cor preta

    # Calculate the box size as a percentage of the image size
    box_percentage = 0.1  # Adjust as needed
    box_size = int(min(output_image.shape[1], output_image.shape[0]) * box_percentage)

    # Calculate the box position both horizontally and vertically
    box_x = (output_image.shape[1] - box_size) // 2
    box_y = (output_image.shape[0] - box_size) // 2

    # Color-code the object box based on the predicted class
    object_box_color = (0, 255, 0) if predicted_class == "Normal" else (255, 0, 0)  # Green for Normal, Red for Cataract

    # Draw a centered object identification box (blue rectangle)
    cv2.rectangle(output_image, (box_x, box_y), (box_x + box_size, box_y + box_size), object_box_color, 2)  # Caixa centralizada

    return output_image

# Crie uma interface Gradio
input_interface = gr.Interface(
    fn=classify_image,
    inputs="image",  # Especifique o tipo de entrada como "image"
    outputs="image",  # Especifique o tipo de saída como "image"
    live=True
)

# Inicie o aplicativo Gradio
input_interface.launch()