import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Nombre del modelo en el repositorio de Hugging Face
model_name = "mkjaramillo/cancer2"

# Cargar el tokenizer y el modelo desde el repositorio de Hugging Face
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Transformación de la imagen
image_transform = transforms.Compose([
    transforms.Resize((50, 50)),
    transforms.ToTensor(),
   
])

def classify_image(image):
    # Cargar la imagen
    image = Image.fromarray(image)
    
    # Preprocesar la imagen
    image = image_transform(image).unsqueeze(0)
    
    # Realizar la inferencia con el modelo
    outputs = model(image)
    
    # Obtener las predicciones
    predictions = torch.argmax(outputs.logits, dim=1)
    
    # Obtener la etiqueta de la predicción
    label = tokenizer.decode(predictions.item())

    # Retornar la etiqueta de la predicción
    return label

# Configurar la interfaz de Gradio
iface = gr.Interface(fn=classify_image,
                     inputs=gr.inputs.Image(label="Imagen de entrada"),
                     outputs="text",
                     title="Clasificador de Imágenes")

# Ejecutar la interfaz de Gradio
iface.launch()