Update app.py

app.py

@@ -12,6 +12,17 @@ from tqdm import tqdm
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_MEAN = [0.485, 0.456, 0.406]
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_STD = [0.229, 0.224, 0.225]
 
+def generate_target_list(data, entity2id):
+    sub = data.loc[(data["datatype_h"] == "image") & (data["datatype_t"] == "id"), ['t']]
+    sub = list(sub['t'])
+    categories = []
+    for item in tqdm(sub):
+        if entity2id[str(int(float(item)))] not in categories:
+            categories.append(entity2id[str(int(float(item)))])
+    # print('categories = {}'.format(categories))
+    # print("No. of target categories = {}".format(len(categories)))
+    return torch.tensor(categories, dtype=torch.long).unsqueeze(-1)
+    
 # Load necessary data and initialize the model
 entity2id = json.load(open('entity2id_subtree.json', 'r'))
 id2entity = {v: k for k, v in entity2id.items()}