Update app.py

app.py

@@ -1,3 +1,20 @@
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch.nn.functional as F
+
+
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+
+def predict(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()
+    predicted_class_idx = logits.argmax(-1).item()
+    label = model.config.id2label[predicted_class_idx].split(",")[0]
+    return {label: float(predicted_class_prob)}
+ 
 import gradio as gr
 
-gr.Interface.load("huggingface/google/vit-base-patch16-224").launch()
+gr.Interface(predict, gr.Image(type="pil"), "label").launch()