Update sroie_inference.py

sroie_inference.py

@@ -46,6 +46,8 @@ def prediction(image):
 
     predictions = outputs.logits.argmax(-1).squeeze().tolist()
     token_boxes = encoding.bbox.squeeze().tolist()
+    probabilities = torch.softmax(outputs.logits, dim=-1)
+    confidence_scores = probabilities.max(-1).values.squeeze().tolist()
 
     inp_ids = encoding.input_ids.squeeze().tolist()
     inp_words = [tokenizer.decode(i) for i in inp_ids]
@@ -55,6 +57,7 @@ def prediction(image):
 
     true_predictions = [id2label[pred] for idx, pred in enumerate(predictions) if not is_subword[idx]]
     true_boxes = [unnormalize_box(box, width, height) for idx, box in enumerate(token_boxes) if not is_subword[idx]]
+    true_confidence_scores = [confidence_scores[idx] for idx, conf in enumerate(confidence_scores) if not is_subword[idx]]
     true_words = []
 
     for id, i in enumerate(inp_words):
@@ -66,16 +69,18 @@ def prediction(image):
     true_predictions = true_predictions[1:-1]
     true_boxes = true_boxes[1:-1]
     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 j != 'others':
-            preds.append(true_predictions[i])
-            l_words.append(true_words[i])
-            bboxes.append(true_boxes[i])
+        if true_confidence_scores[i] < 0.9:  #####################################àà
+            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):