In the parameters I have put default values for the size (=1) and thickness (=2) of the label. If the user wants to change these values, when the function he wants is called, he will put the values he wants after the bbox, first for the size and then for the thickness.

Author: nefelipar

bbox_visualizer/bbox_visualizer.py

@@ -1,5 +1,7 @@
 import cv2
 
+font = cv2.FONT_HERSHEY_SIMPLEX
+
 
 def draw_rectangle(img,
                    bbox,
@@ -47,6 +49,8 @@ def draw_rectangle(img,
 def add_label(img,
               label,
               bbox,
+              size=1,
+              thickness=2,
               draw_bg=True,
               text_bg_color=(255, 255, 255),
               text_color=(0, 0, 0),
@@ -61,6 +65,10 @@ def add_label(img,
         the text (label) to be written
     bbox : list
         a list containing x_min, y_min, x_max and y_max of the rectangle positions
+    size : int, optional
+        size of the label, by default 1
+    thickness : int, optional
+        thickness of the label, by default 2
     draw_bg : bool, optional
         if True, draws the background of the text, else just the text is written, by default True
     text_bg_color : tuple, optional
@@ -76,30 +84,29 @@ def add_label(img,
         the image with the label written
     """
 
-    text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][0]
+    (label_width, label_height), baseline = cv2.getTextSize(label, font, size, thickness)
 
     if top:
-        label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] - 30]
+        label_bg = [bbox[0], bbox[1], bbox[0] + label_width, bbox[1] - label_height - (15 * size)]
         if draw_bg:
             cv2.rectangle(img, (label_bg[0], label_bg[1]),
                           (label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
-        cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5),
-                    cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
 
+        cv2.putText(img, label, (bbox[0] + 5, bbox[1] - (15 * size)), font, size, text_color, thickness)
     else:
-        label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] + 30]
+        label_bg = [bbox[0], bbox[1], bbox[0] + label_width, bbox[1] + label_height + (15 * size)]
         if draw_bg:
             cv2.rectangle(img, (label_bg[0], label_bg[1]),
                           (label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
-        cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5 + 30),
-                    cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
-
+        cv2.putText(img, label, (bbox[0] + 5, bbox[1] + (16 * size) + (4 * thickness)), font, size, text_color, thickness)
     return img
 
 
 def add_T_label(img,
                 label,
                 bbox,
+                size=1,
+                thickness=2,
                 draw_bg=True,
                 text_bg_color=(255, 255, 255),
                 text_color=(0, 0, 0)):
@@ -113,6 +120,10 @@ def add_T_label(img,
         the text (label) to be written
     bbox : list
         a list containing x_min, y_min, x_max and y_max of the rectangle positions
+    size : int, optional
+        size of the label, by default 1
+    thickness : int, optional
+        thickness of the label, by default 2
     draw_bg : bool, optional
         if True, draws the background of the text, else just the text is written, by default True
     text_bg_color : tuple, optional
@@ -126,35 +137,36 @@ def add_T_label(img,
         the image with the T label drawn/written
     """
 
-    text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][0]
-    text_height = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][1]
-
+    (label_width, label_height), baseline = cv2.getTextSize(label, font, size, thickness)
     # draw vertical line
     x_center = (bbox[0] + bbox[2]) // 2
     y_top = bbox[1] - 50
     cv2.line(img, (x_center, bbox[1]), (x_center, y_top), text_bg_color, 3)
 
     # draw rectangle with label
     y_bottom = y_top
-    y_top = y_bottom - text_height - 5
-    x_left = x_center - (text_width // 2) - 5
-    x_right = x_center + (text_width // 2) + 5
+    y_top = y_bottom - label_height - 5
+    x_left = x_center - (label_width // 2) - 5
+    x_right = x_center + (label_width // 2) + 5
     if draw_bg:
-        cv2.rectangle(img, (x_left, y_top - 3), (x_right, y_bottom),
+        cv2.rectangle(img, (x_left, y_top - 30), (x_right, y_bottom),
                       text_bg_color, -1)
-    cv2.putText(img, label, (x_left + 5, y_bottom - 7),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
+    cv2.putText(img, label, (x_left + 5, y_bottom - (8 * size)),
+                font, size, text_color, thickness)
 
     return img
 
 
 def draw_flag_with_label(img,
                          label,
                          bbox,
+                         size=1,
+                         thickness=2,
                          write_label=True,
                          line_color=(255, 255, 255),
                          text_bg_color=(255, 255, 255),
-                         text_color=(0, 0, 0)):
+                         text_color=(0, 0, 0),
+                         ):
     """draws a pole from the middle of the object that is to be labeled and adds the label to the flag
 
     Parameters
@@ -165,6 +177,10 @@ def draw_flag_with_label(img,
         label that is written inside the flag
     bbox : list
         a list containing x_min, y_min, x_max and y_max of the rectangle positions
+    size : int, optional
+        size of the label, by default 1
+    thickness : int, optional
+        thickness of the label, by default 2
     write_label : bool, optional
         if True, writes the label, otherwise, it's just a vertical line, by default True
     line_color : tuple, optional
@@ -181,6 +197,7 @@ def draw_flag_with_label(img,
     """
 
     # draw vertical line
+    (label_width, label_height), baseline = cv2.getTextSize(label, font, size, thickness)
 
     x_center = (bbox[0] + bbox[2]) // 2
     y_bottom = int((bbox[1] * .75 + bbox[3] * .25))
@@ -194,17 +211,14 @@ def draw_flag_with_label(img,
     # write label
 
     if write_label:
-        text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1,
-                                     2)[0][0]
         label_bg = [
-            start_point[0], start_point[1], start_point[0] + text_width,
-            start_point[1] + 30
+            start_point[0], start_point[1], start_point[0] + label_width,
+            start_point[1] - label_height - (10 * size)
         ]
         cv2.rectangle(img, (label_bg[0], label_bg[1]),
                       (label_bg[2] + 5, label_bg[3]), text_bg_color, -1)
-        cv2.putText(img, label, (start_point[0] + 5, start_point[1] - 5 + 30),
-                    cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
-
+        cv2.putText(img, label, (start_point[0] + 7, start_point[1] - (13 * size) + 20),
+                    font, size, text_color, thickness)
     return img