AppCancer / app.py
mkjaramillo's picture
Update app.py
2332cfc
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()