complete app

app.py

@@ -1,7 +1,50 @@
 import gradio as gr
+from transformers import ViTForImageClassification
+from torchvision import transforms
+import os 
+import numpy as np
+import json
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+class WordVocabulary:
+    def __init__(self, records=None, word2idx_path=None):
+        if word2idx_path is not None:
+            self.records = []
+            self.word2idx = json.load(open(word2idx_path, "r"))
+            self.words = list(self.word2idx.keys())
+            return
+
+
+    def build_vocabulary(self):
+        words = set()
+        for r in self.records:
+            words.update([w.strip() for w in r['text'].split(",")])
+        self.words = sorted(list(words))
+        self.word2idx = {w: idx for (idx, w) in enumerate(self.words)}
+
+
+vocabulary = WordVocabulary(word2idx_path="word2idx.json")
+model = ViTForImageClassification.from_pretrained("Inf009/food1024_vit_focal_mixup", problem_type="multi_label_classification", num_labels=len(vocabulary))
+test_transforms = transforms.Compose(
+    [
+        transforms.Resize((256, 256)),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+    ]
+)
+
+def multi_label_predict(img, threshold=0.5):
+    img_transformed = test_transforms(img)
+    outputs = model(img_transformed.unsqueeze(0)).logits.squeeze(0).sigmoid().detach().numpy()
+    indices = np.where(outputs > threshold)[0]
+    indices = sorted(indices, key=lambda x: outputs[x], reverse=True)
+    predict_tags = [vocabulary[idx] for idx in indices]
+    return predict_tags
+
+demo_image_path = "images"
+images = [f for f in os.listdir(demo_image_path) if f.endswith(".jpg")][:10]
+images = [os.path.join(demo_image_path, file) for file in images]
+examples = [[image, 0.5] for image in images]
+iface = gr.Interface(fn=multi_label_predict, inputs=[gr.inputs.Image(type="pil"), gr.inputs.Number(default=0.5)], 
+             examples=examples, outputs="text")
 iface.launch()

word2idx.json

@@ -0,0 +1,105 @@
+{
+    "French beans": 0,
+    "almond": 1,
+    "apple": 2,
+    "apricot": 3,
+    "asparagus": 4,
+    "avocado": 5,
+    "bamboo shoots": 6,
+    "banana": 7,
+    "bean sprouts": 8,
+    "biscuit": 9,
+    "blueberry": 10,
+    "bread": 11,
+    "broccoli": 12,
+    "cabbage": 13,
+    "cake": 14,
+    "candy": 15,
+    "carrot": 16,
+    "cashew": 17,
+    "cauliflower": 18,
+    "celery stick": 19,
+    "cheese butter": 20,
+    "cherry": 21,
+    "chicken duck": 22,
+    "chocolate": 23,
+    "cilantro mint": 24,
+    "coffee": 25,
+    "corn": 26,
+    "crab": 27,
+    "cucumber": 28,
+    "date": 29,
+    "dried cranberries": 30,
+    "egg": 31,
+    "egg tart": 32,
+    "eggplant": 33,
+    "enoki mushroom": 34,
+    "fig": 35,
+    "fish": 36,
+    "french fries": 37,
+    "fried meat": 38,
+    "garlic": 39,
+    "ginger": 40,
+    "grape": 41,
+    "green beans": 42,
+    "hamburg": 43,
+    "hanamaki baozi": 44,
+    "ice cream": 45,
+    "juice": 46,
+    "kelp": 47,
+    "king oyster mushroom": 48,
+    "kiwi": 49,
+    "lamb": 50,
+    "lemon": 51,
+    "lettuce": 52,
+    "mango": 53,
+    "melon": 54,
+    "milk": 55,
+    "milkshake": 56,
+    "noodles": 57,
+    "okra": 58,
+    "olives": 59,
+    "onion": 60,
+    "orange": 61,
+    "other ingredients": 62,
+    "oyster mushroom": 63,
+    "pasta": 64,
+    "peach": 65,
+    "peanut": 66,
+    "pear": 67,
+    "pepper": 68,
+    "pie": 69,
+    "pineapple": 70,
+    "pizza": 71,
+    "popcorn": 72,
+    "pork": 73,
+    "potato": 74,
+    "pudding": 75,
+    "pumpkin": 76,
+    "rape": 77,
+    "raspberry": 78,
+    "red beans": 79,
+    "rice": 80,
+    "salad": 81,
+    "sauce": 82,
+    "sausage": 83,
+    "seaweed": 84,
+    "shellfish": 85,
+    "shiitake": 86,
+    "shrimp": 87,
+    "snow peas": 88,
+    "soup": 89,
+    "soy": 90,
+    "spring onion": 91,
+    "steak": 92,
+    "strawberry": 93,
+    "tea": 94,
+    "tofu": 95,
+    "tomato": 96,
+    "walnut": 97,
+    "watermelon": 98,
+    "white button mushroom": 99,
+    "white radish": 100,
+    "wine": 101,
+    "wonton dumplings": 102
+}