ClassCat commited on
Commit
dfe0d6f
1 Parent(s): 8a2a602

add app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from transformers import ViTFeatureExtractor, ViTForImageClassification
3
+
4
+ import torch
5
+ import gradio as gr
6
+
7
+ from PIL import Image
8
+
9
+ feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
10
+ model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
11
+
12
+ def classify_image(image):
13
+
14
+ with torch.no_grad():
15
+ model.eval()
16
+
17
+ inputs = feature_extractor(images=image, return_tensors="pt")
18
+ outputs = model(**inputs)
19
+
20
+ logits = outputs.logits
21
+ prob = torch.nn.functional.softmax(logits, dim=1)
22
+
23
+ top10_prob, top10_indices = torch.topk(prob, 10)
24
+ top10_confidences = {}
25
+ for i in range(10):
26
+ top10_confidences[model.config.id2label[int(top10_indices[0][i])]] = float(top10_prob[0][i])
27
+
28
+ return top10_confidences #confidences
29
+
30
+
31
+ with gr.Blocks(title="ViT ImageNet Classification - ClassCat",
32
+ css=".gradio-container {background:mintcream;}"
33
+ ) as demo:
34
+ gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">ViT - ImageNet Classification</div>""")
35
+
36
+ with gr.Row():
37
+ input_image = gr.Image(type="pil", image_mode="RGB", shape=(224, 224))
38
+ output_label=gr.Label(label="Probabilities", num_top_classes=3)
39
+
40
+ send_btn = gr.Button("Infer")
41
+ send_btn.click(fn=classify_image, inputs=input_image, outputs=output_label)
42
+
43
+ #demo.queue(concurrency_count=3)
44
+ demo.launch(debug=True)