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

# Carregue o extrator de recursos e o modelo
extractor = AutoFeatureExtractor.from_pretrained("SerdarHelli/ThyroidTumorClassificationModel")
model = AutoModelForImageClassification.from_pretrained("SerdarHelli/ThyroidTumorClassificationModel")

# Função para classificar a imagem
def classify_image(image):
    # Pré-processa a imagem usando o extrator
    inputs = extractor(images=image, return_tensors="pt")
    
    # Passa a imagem pelo modelo
    outputs = model(**inputs)
    
    # 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()
    
    # Rótulos de classe personalizados com base no seu modelo
    class_labels = ["Sem Tumor", "Tumor"]
    
    # Rótulo da classe prevista
    predicted_label = class_labels[predicted_class]
    
    # 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
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(),
    outputs=gr.outputs.HTML(),  # Saída formatada com HTML
    title="Classificador de Tumor da Tireoide",
    description="Faça o upload de uma imagem de um tumor da tireoide para classificação.",
)

# Inicie a interface Gradio
iface.launch()