Update cord_inference.py

cord_inference.py

@@ -30,6 +30,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]
@@ -39,6 +41,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):
@@ -50,6 +53,7 @@ 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 = []