From dc4602c75a48780fc9a500a40afe5861311b303b Mon Sep 17 00:00:00 2001
From: jaycecd <33342388+jaycecd@users.noreply.github.com>
Date: Sun, 1 May 2022 20:02:54 +0800
Subject: [PATCH] display probability scores on visualization  (#5854)

* display probability scores on visualization

* display probability scores on visualization
---
 tools/infer/utility.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index 0bbc3b1d..18494e37 100644
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -429,6 +429,7 @@ def draw_ocr_box_txt(image,
                      boxes,
                      txts,
                      scores=None,
+                     show_score=False,
                      drop_score=0.5,
                      font_path="./doc/simfang.ttf"):
     h, w = image.height, image.width
@@ -440,7 +441,7 @@ def draw_ocr_box_txt(image,
     random.seed(0)
     draw_left = ImageDraw.Draw(img_left)
     draw_right = ImageDraw.Draw(img_right)
-    for idx, (box, txt) in enumerate(zip(boxes, txts)):
+    for idx, (box, txt, score) in enumerate(zip(boxes, txts, scores)):
         if scores is not None and scores[idx] < drop_score:
             continue
         color = (random.randint(0, 255), random.randint(0, 255),
@@ -456,6 +457,8 @@ def draw_ocr_box_txt(image,
             1])**2)
         box_width = math.sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][
             1])**2)
+        if show_score:
+            txt = txt + ':' + str(score)
         if box_height > 2 * box_width:
             font_size = max(int(box_width * 0.9), 10)
             font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
-- 
GitLab