yangswei commited on
Commit
1567d28
·
verified ·
1 Parent(s): 0b905d3

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -0
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import ViTFeatureExtractor, ViTForImageClassification
3
+ from PIL import Image
4
+ import requests
5
+ url = 'https://potatorolls.com/wp-content/uploads/2020/10/Basic-Hot-Dogs-960x640.jpg'
6
+ hotdog_image = Image.open(requests.get(url, stream=True).raw)
7
+
8
+ def snacks_classifier(input_image):
9
+
10
+ # Init model, transforms
11
+ processor = ViTImageProcessor.from_pretrained('yangswei/snacks_classification')
12
+ model = ViTForImageClassification.from_pretrained('yangswei/snacks_classification')
13
+
14
+ # inputs & outputs
15
+ inputs = processor(images=input_image, return_tensors="pt")
16
+ outputs = model(**inputs).logits.softmax(1)
17
+ labels = model.config.id2label
18
+ confidences = {labels[i]: outputs[0][i].item() for i in range(len(labels))}
19
+
20
+ return confidences
21
+
22
+ with gr.Blocks(theme=gr.themes.Base()) as demo:
23
+ gr.Interface(fn=snacks_classifier, inputs="image", outputs=gr.Label(num_top_classes=20, label="Prediction"),
24
+ examples=[hotdog_image])
25
+
26
+ demo.launch()