app

app.py

@@ -18,7 +18,7 @@ def mlm(image, text):
             #'def': f'{cls_name} is a  . [SEP]',
         }
         ans = list()
-        res = defaultdict(list)
+        res = defaultdict()
         device = 'cpu'
         for k, v in questions_dict.items():
             predicted_tokens = []
@@ -36,15 +36,18 @@ def mlm(image, text):
             with torch.no_grad():
                 predictions = model(tokens_tensor, segments_tensors)
             
-            _, predicted_index = torch.topk(predictions[0][0][masked_index], 2)#.item()
+            _, predicted_index = torch.topk(predictions[0][0][masked_index], 1)#.item()
             predicted_index = predicted_index.detach().cpu().numpy()
             #print(predicted_index)
             for idx in predicted_index:
                 predicted_tokens.append(tokenizer.convert_ids_to_tokens([idx])[0])
-            for i in range(2):
-                res[text][k].append(predicted_tokens)
+            # for i in range(1):
+            #     res[text][k].append(predicted_tokens)
+            res[k] = predicted_tokens[0]
+        color, shape, loc = res['color'], res['shape'], res['location']
+        ans = f'{color} color, {shape} shape, cat at {loc}'
 
-        return image, res
+        return image, ans
 
 def to_black(image, text):
     output = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)