app.py

import gradio as gr
from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import torch

# Cargar el modelo y el extractor de características
model = ViTForImageClassification.from_pretrained("akahana/vit-base-cats-vs-dogs")
feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")

# Función de predicción
def classify_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    predicted_class = model.config.id2label[predicted_class_idx]
    return predicted_class

# Crear la interfaz de Gradio
interface = gr.Interface(fn=classify_image, inputs=gr.Image(type="pil"), outputs="text")
interface.launch()