Create app.py

app.py

@@ -0,0 +1,26 @@
+import gradio as gr
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import requests
+url = 'https://potatorolls.com/wp-content/uploads/2020/10/Basic-Hot-Dogs-960x640.jpg'
+hotdog_image = Image.open(requests.get(url, stream=True).raw)
+
+def snacks_classifier(input_image):
+
+    # Init model, transforms
+    processor = ViTImageProcessor.from_pretrained('yangswei/snacks_classification')
+    model = ViTForImageClassification.from_pretrained('yangswei/snacks_classification')
+
+    # inputs & outputs
+    inputs = processor(images=input_image, return_tensors="pt")
+    outputs = model(**inputs).logits.softmax(1)
+    labels = model.config.id2label
+    confidences = {labels[i]: outputs[0][i].item() for i in range(len(labels))}
+
+    return confidences
+
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+    gr.Interface(fn=snacks_classifier, inputs="image", outputs=gr.Label(num_top_classes=20, label="Prediction"),
+                 examples=[hotdog_image])
+
+demo.launch()