Update sroie_inference.py

sroie_inference.py

@@ -71,24 +71,16 @@ def prediction(image):
     true_words = true_words[1:-1]
     true_confidence_scores = true_confidence_scores[1:-1]
 
-    preds = []
-    l_words = []
-    bboxes = []
-
-    for i, j in enumerate(true_predictions):
-        if true_confidence_scores[i] < 0.9:  #####################################àà
+    for i, j in enumerate(true_confidence_scores):
+        if j < 0.9 or :  #####################################
             true_predictions[i] = "O"
-        preds.append(true_predictions[i])
-        l_words.append(true_words[i])
-        bboxes.append(true_boxes[i])
 
     d = {}
-    for id, i in enumerate(preds):
+    for id, i in enumerate(true_predictions):
         if i not in d.keys():
-            d[i] = l_words[id]
+            d[i] = true_words[d]
         else:
-            d[i] = d[i] + ", " + l_words[id]
-
+            d[i] = d[i] + ", " + true_words[id]
     d = {k: v.strip() for (k, v) in d.items()}
 
     keys_to_pop = []
@@ -101,6 +93,9 @@ def prediction(image):
     if "B-TOTAL" in d: d.pop("B-TOTAL")
     for k in keys_to_pop: d.pop(k)
 
+    for k in d.keys():
+        k = k[2:]
+
     blur_boxes = []
     for prediction, box in zip(preds, bboxes):
         if prediction != 'O' and prediction[2:] != 'TOTAL':