optimize box visualize

common/utils.py

@@ -5,6 +5,7 @@
 import os
 import numpy as np
 import time
+from enum import Enum
 import cv2, colorsys
 from PIL import Image
 from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
@@ -116,7 +117,12 @@ def get_dataset(annotation_file, shuffle=True):
 
     return lines
 
-def draw_label(image, text, color, coords):
+labelType = Enum('labelType', ('LABEL_TOP_OUTSIDE',
+                               'LABEL_BOTTOM_OUTSIDE',
+                               'LABEL_TOP_INSIDE',
+                               'LABEL_BOTTOM_INSIDE',))
+
+def draw_label(image, text, color, coords, label_type=labelType.LABEL_TOP_OUTSIDE):
     font = cv2.FONT_HERSHEY_PLAIN
     font_scale = 1.
     (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]
@@ -127,11 +133,18 @@ def draw_label(image, text, color, coords):
 
     (x, y) = coords
 
-    cv2.rectangle(image, (x, y), (x + rect_width, y - rect_height), color, cv2.FILLED)
-    cv2.putText(image, text, (x + padding, y - text_height + padding), font,
-                fontScale=font_scale,
-                color=(255, 255, 255),
-                lineType=cv2.LINE_AA)
+    if label_type == labelType.LABEL_TOP_OUTSIDE or label_type == labelType.LABEL_BOTTOM_INSIDE:
+        cv2.rectangle(image, (x, y), (x + rect_width, y - rect_height), color, cv2.FILLED)
+        cv2.putText(image, text, (x + padding, y - text_height + padding), font,
+                    fontScale=font_scale,
+                    color=(255, 255, 255),
+                    lineType=cv2.LINE_AA)
+    else: # LABEL_BOTTOM_OUTSIDE or LABEL_TOP_INSIDE
+        cv2.rectangle(image, (x, y), (x + rect_width, y + rect_height), color, cv2.FILLED)
+        cv2.putText(image, text, (x + padding, y + text_height + padding), font,
+                    fontScale=font_scale,
+                    color=(255, 255, 255),
+                    lineType=cv2.LINE_AA)
 
     return image
 
@@ -156,8 +169,20 @@ def draw_boxes(image, boxes, classes, scores, class_names, colors, show_score=Tr
             color = (0,0,0)
         else:
             color = colors[cls]
+
+        # choose label type according to box size
+        if ymin > 20:
+            label_coords = (xmin, ymin)
+            label_type = label_type=labelType.LABEL_TOP_OUTSIDE
+        elif ymin <= 20 and ymax <= image.shape[0] - 20:
+            label_coords = (xmin, ymax)
+            label_type = label_type=labelType.LABEL_BOTTOM_OUTSIDE
+        elif ymax > image.shape[0] - 20:
+            label_coords = (xmin, ymin)
+            label_type = label_type=labelType.LABEL_TOP_INSIDE
+
         cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 1, cv2.LINE_AA)
-        image = draw_label(image, label, color, (xmin, ymin))
+        image = draw_label(image, label, color, label_coords, label_type)
 
     return image
 

tools/misc/augment_test.py

@@ -7,11 +7,12 @@
 import numpy as np
 from PIL import Image
 import cv2
+from tqdm import tqdm
 
 sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', '..'))
 from yolo3.data import get_ground_truth_data
-from common.utils import get_dataset, get_classes, draw_label
-from common.data_utils import random_mosaic_augment, random_mosaic_augment_v5, random_cutmix_augment
+from common.utils import get_dataset, get_classes, draw_label, labelType
+from common.data_utils import random_mosaic_augment, random_mosaic_augment_v5, random_cutmix_augment, denormalize_image
 
 
 def draw_boxes(images, boxes, class_names, output_path):
@@ -22,17 +23,29 @@ def draw_boxes(images, boxes, class_names, output_path):
             if (boxes[i][j][:4] == 0).all():
                 continue
 
-            x_min = int(boxes[i][j][0])
-            y_min = int(boxes[i][j][1])
-            x_max = int(boxes[i][j][2])
-            y_max = int(boxes[i][j][3])
+            xmin = int(boxes[i][j][0])
+            ymin = int(boxes[i][j][1])
+            xmax = int(boxes[i][j][2])
+            ymax = int(boxes[i][j][3])
             cls = int(boxes[i][j][4])
 
             class_name = class_names[cls]
             color = (255,0,0)  #Red box
 
-            cv2.rectangle(img, (x_min, y_min), (x_max, y_max), color, 2)
-            img = draw_label(img, class_name, color, (x_min, y_min))
+            # choose label type according to box size
+            if ymin > 20:
+                label_coords = (xmin, ymin)
+                label_type = label_type=labelType.LABEL_TOP_OUTSIDE
+            elif ymin <= 20 and ymax <= img.shape[0] - 20:
+                label_coords = (xmin, ymax)
+                label_type = label_type=labelType.LABEL_BOTTOM_OUTSIDE
+            elif ymax > img.shape[0] - 20:
+                label_coords = (xmin, ymin)
+                label_type = label_type=labelType.LABEL_TOP_INSIDE
+
+            cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 1, cv2.LINE_AA)
+            img = draw_label(img, class_name, color, label_coords, label_type)
+
         image = Image.fromarray(img)
         image.save(os.path.join(output_path, str(i)+".jpg"))
 
@@ -57,15 +70,18 @@ def main():
 
     image_data = []
     boxes_data = []
+
+    pbar = tqdm(total=args.batch_size, desc='Generate augment image')
     for i in range(args.batch_size):
+        pbar.update(1)
         annotation_line = annotation_lines[i]
         image, boxes = get_ground_truth_data(annotation_line, input_shape=model_input_shape, augment=True)
-        #un-normalize image
-        image = image*255.0
-        image = image.astype(np.uint8)
+        #denormalize image
+        image = denormalize_image(image)
 
         image_data.append(image)
         boxes_data.append(boxes)
+    pbar.close()
     image_data = np.array(image_data)
     boxes_data = np.array(boxes_data)