import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import torch
import datetime
from PIL import Image
import torchvision.transforms as transforms

# Carregue o extrator de recursos e o modelo
extractor = AutoFeatureExtractor.from_pretrained("AhmadHakami/alzheimer-image-classification-google-vit-base-patch16")
model = AutoModelForImageClassification.from_pretrained("AhmadHakami/alzheimer-image-classification-google-vit-base-patch16")

# Transformações para pré-processamento
preprocess = transforms.Compose([
    transforms.ToPILImage(),  # Converte para objeto PIL.Image
    transforms.Resize((224, 224)),  # Redimensiona para o tamanho de entrada do modelo (224x224)
    transforms.ToTensor(),  # Converte para tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalização
])

# Função para classificar a imagem
def classify_image(image):
    # Aplica o pré-processamento na imagem
    image = preprocess(image)
    
    # Adiciona uma dimensão extra para criar um lote (batch) de tamanho 1
    image = image.unsqueeze(0)
    
    # Passa a imagem pelo modelo
    outputs = model(image)
    
    # Obtém as probabilidades das classes
    logits = outputs.logits
    
    # Calcula as probabilidades finais usando o softmax
    probabilities = torch.softmax(logits, dim=1)
    
    # Obtém a classe com a maior probabilidade
    predicted_class = torch.argmax(probabilities, dim=1).item()
    
    # Mapeamento de classe ID para rótulo
    id2label = {
        "0": "Mild_Demented",
        "1": "Moderate_Demented",
        "2": "Non_Demented",
        "3": "Very_Mild_Demented"
    }
    
    # Rótulo da classe prevista
    predicted_label = id2label.get(str(predicted_class), "Desconhecido")
    
    # Obtém a data e hora atual
    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Formate a saída em HTML com data e hora
    result_html = f"""
    <h2>Resultado da Classificação</h2>
    <p><strong>Classe Predita:</strong> {predicted_label}</p>
    <p><strong>Data e Hora:</strong> {current_time}</p>
    """
    
    # Retorna o resultado formatado em HTML
    return result_html

# Crie uma interface Gradio com instruções de uso
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(type="numpy", label="Carregar uma imagem"),
    outputs=gr.outputs.HTML(),  # Saída formatada com HTML
    title="Classificador de Imagem ViT para Demência com Pré-processamento",
    description="""<h1>Classificador de Imagem ViT para Demência</h1>
<p>Esta aplicação Gradio permite classificar imagens relacionadas à demência usando um modelo Vision Transformer (ViT).</p>

<h2>Instruções de Uso</h2>

<ol>
  <li>Clique no botão "Escolher Arquivo" para fazer o upload de uma imagem relacionada à demência.</li>
  <li>Aguarde alguns segundos enquanto o modelo processa a imagem.</li>
  <li>O resultado mostrará a classe prevista relacionada à demência.</li>
  <li>A data e a hora da classificação também serão exibidas.</li>
</ol>

<p>Este modelo pode identificar classes de demência, como "Mild_Demented", "Moderate_Demented", "Non_Demented" e "Very_Mild_Demented".</p>

<p>Por favor, note que este é um modelo de demonstração e os resultados podem não ser precisos para fins clínicos.</p>"""
)

# Inicie a interface Gradio
iface.launch()