mohd43 commited on
Commit
dc69039
1 Parent(s): a2c55f0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import ViTImageProcessor, ViTForImageClassification
2
+ import gradio as gr
3
+ from PIL import Image
4
+ import requests
5
+
6
+
7
+ processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
8
+ model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
9
+
10
+
11
+ def predict(image) :
12
+ inputs = processor(images=image, return_tensors="pt")
13
+ outputs = model(**inputs)
14
+ logits = outputs.logits
15
+ # model predicts one of the 1000 ImageNet classes
16
+ predicted_class_idx = logits.argmax(-1).item()
17
+ return model.config.id2label[predicted_class_idx]
18
+
19
+
20
+
21
+
22
+ gradio_app = gr.Interface(
23
+ predict,
24
+ inputs=gr.Image(label="Select image for classification", sources=['upload', 'webcam'], type="pil"),
25
+ outputs=gr.Textbox(),
26
+ title="Image Classification",
27
+ live=True,
28
+ allow_flagging="never",
29
+ )
30
+
31
+ gradio_app.launch()