Add application file

app.py

@@ -1,7 +1,21 @@
 import gradio as gr
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
+# 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")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# 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()