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 con título y descripción
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Clasificador de Gatos vs Perros",
    description="Sube una imagen de un gato o un perro y este modelo basado en ViT te dirá qué animal aparece en la imagen."
)

interface.launch()