Upload 2 files

app.py

@@ -0,0 +1,29 @@
+import gradio as gr
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+
+# Load pre-trained model and feature extractor
+model_name = "google/vit-base-patch16-224"
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
+
+# Define the prediction function
+def classify_image(img):
+    inputs = feature_extractor(images=img, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    predicted_label = model.config.id2label[predicted_class_idx]
+    return predicted_label
+
+# Build the Gradio interface
+interface = gr.Interface(fn=classify_image,
+                         inputs=gr.Image(type="pil"),
+                         outputs="text",
+                         title="Image Classification with ViT",
+                         description="Upload an image and classify it using Vision Transformer (ViT)")
+
+# Launch the app
+interface.launch()

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+transformers
+torch
+