Merge pull request #12 from THEFASHIONGEEK/THEFASHIONGEEK-patch-2

Added confidence value - Verified and merged.
Tessellate-Imaging · Mar 19, 2020 · e006218 · e006218
2 parents 8395c71 + ce1301b
commit e006218
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/6_cornernet_lite/lib/core/vis_utils.py b/6_cornernet_lite/lib/core/vis_utils.py
@@ -3,7 +3,6 @@
 
 def draw_bboxes(image, bboxes, font_size=0.5, thresh=0.5, colors=None):
     """Draws bounding boxes on an image.
-
     Args:
         image: An image in OpenCV format
         bboxes: A dictionary representing bounding boxes of different object
@@ -16,15 +15,14 @@ def draw_bboxes(image, bboxes, font_size=0.5, thresh=0.5, colors=None):
             will be drawn.
         colors: (Optional) Color of bounding boxes for each category. If it is
             not provided, this function will use random color for each category.
-
     Returns:
         An image with bounding boxes.
     """
 
     image = image.copy()
     for cat_name in bboxes:
         keep_inds = bboxes[cat_name][:, -1] > thresh
-        cat_size  = cv2.getTextSize(cat_name, cv2.FONT_HERSHEY_SIMPLEX, font_size, 2)[0]
+        cat_size  = cv2.getTextSize(cat_name +' : %.2f' % bboxes[cat_name][:, -1], cv2.FONT_HERSHEY_SIMPLEX, font_size, 2)[0]
 
         if colors is None:
             color = np.random.random((3, )) * 0.6 + 0.4
@@ -40,7 +38,7 @@ def draw_bboxes(image, bboxes, font_size=0.5, thresh=0.5, colors=None):
                     (bbox[0] + cat_size[0], bbox[1] + cat_size[1] + 2),
                     color, -1
                 )
-                cv2.putText(image, cat_name,
+                cv2.putText(image, cat_name +' : %.2f' % bboxes[cat_name][:, -1],
                     (bbox[0], bbox[1] + cat_size[1] + 2),
                     cv2.FONT_HERSHEY_SIMPLEX, font_size, (0, 0, 0), thickness=1
                 )
@@ -50,7 +48,7 @@ def draw_bboxes(image, bboxes, font_size=0.5, thresh=0.5, colors=None):
                     (bbox[0] + cat_size[0], bbox[1] - 2),
                     color, -1
                 )
-                cv2.putText(image, cat_name,
+                cv2.putText(image, cat_name +' : %.2f' % bboxes[cat_name][:, -1],
                     (bbox[0], bbox[1] - 2),
                     cv2.FONT_HERSHEY_SIMPLEX, font_size, (0, 0, 0), thickness=1
                 )