example-apps/PDD/pcb-defect-detection/libs/box_utils/draw_box_in_img.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import, print_function, division
   4
   5 import numpy as np
   6
   7 from PIL import Image, ImageDraw, ImageFont
   8 import cv2
   9
  10 from libs.configs import cfgs
  11 from libs.label_name_dict.label_dict import LABEl_NAME_MAP
  12 NOT_DRAW_BOXES = 0
  13 ONLY_DRAW_BOXES = -1
  14 ONLY_DRAW_BOXES_WITH_SCORES = -2
  15
  16 STANDARD_COLORS = [
  17     'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque',
  18     'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite',
  19     'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan',
  20     'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange',
  21     'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet',
  22     'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite',
  23     'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod',
  24     'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki',
  25     'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue',
  26     'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey',
  27     'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue',
  28     'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime',
  29     'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid',
  30     'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen',
  31     'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin',
  32     'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed',
  33     'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed',
  34     'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple',
  35     'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown',
  36     'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue',
  37     'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow',
  38     'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White',
  39     'WhiteSmoke', 'Yellow', 'YellowGreen', 'LightBlue', 'LightGreen'
  40 ]
  41 FONT = ImageFont.load_default()
  42
  43
  44 def draw_a_rectangel_in_img(draw_obj, box, color, width):
  45     '''
  46     use draw lines to draw rectangle. since the draw_rectangle func can not modify the width of rectangle
  47     :param draw_obj:
  48     :param box: [x1, y1, x2, y2]
  49     :return:
  50     '''
  51     x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
  52     top_left, top_right = (x1, y1), (x2, y1)
  53     bottom_left, bottom_right = (x1, y2), (x2, y2)
  54
  55     draw_obj.line(xy=[top_left, top_right],
  56                   fill=color,
  57                   width=width)
  58     draw_obj.line(xy=[top_left, bottom_left],
  59                   fill=color,
  60                   width=width)
  61     draw_obj.line(xy=[bottom_left, bottom_right],
  62                   fill=color,
  63                   width=width)
  64     draw_obj.line(xy=[top_right, bottom_right],
  65                   fill=color,
  66                   width=width)
  67
  68
  69 def only_draw_scores(draw_obj, box, score, color):
  70
  71     x, y = box[0], box[1]
  72     draw_obj.rectangle(xy=[x, y-10, x+100, y],
  73                        fill=color)
  74     draw_obj.text(xy=(x, y-10),
  75                   text="obj:" +str(round(score, 2)),
  76                   fill='white',
  77                   font=FONT)
  78
  79
  80 def draw_label_with_scores(draw_obj, box, label, score, color):
  81     x, y = box[0], box[1]
  82     draw_obj.rectangle(xy=[x, y-10, x+100, y],
  83                        fill=color)
  84
  85     txt = LABEl_NAME_MAP[label] + ':' + str(round(score, 2))
  86     draw_obj.text(xy=(x, y-10),
  87                   text=txt,
  88                   fill='white',
  89                   font=FONT)
  90
  91
  92 def draw_boxes_with_label_and_scores(img_array, boxes, labels, scores):
  93
  94     img_array = img_array + np.array(cfgs.PIXEL_MEAN)
  95     img_array.astype(np.float32)
  96     boxes = boxes.astype(np.int64)
  97     labels = labels.astype(np.int32)
  98     img_array = np.array(img_array * 255 / np.max(img_array), dtype=np.uint8)
  99
 100     img_obj = Image.fromarray(img_array)
 101     raw_img_obj = img_obj.copy()
 102
 103     draw_obj = ImageDraw.Draw(img_obj)
 104     num_of_objs = 0
 105     for box, a_label, a_score in zip(boxes, labels, scores):
 106
 107         if a_label != NOT_DRAW_BOXES:
 108             num_of_objs += 1
 109             #draw_a_rectangel_in_img(draw_obj, box, color=STANDARD_COLORS[a_label], width=3)
 110             draw_a_rectangel_in_img(draw_obj, box, color='Red', width=3)
 111             if a_label == ONLY_DRAW_BOXES:  # -1
 112                 continue
 113             elif a_label == ONLY_DRAW_BOXES_WITH_SCORES:  # -2
 114                  only_draw_scores(draw_obj, box, a_score, color='Yellow')
 115                  continue
 116             else:
 117                 draw_label_with_scores(draw_obj, box, a_label, a_score, color='DeepSkyBlue')
 118
 119     out_img_obj = Image.blend(raw_img_obj, img_obj, alpha=0.7)
 120
 121     return np.array(out_img_obj)
 122
 123
 124 if __name__ == '__main__':
 125     img_array = cv2.imread("/home/yjr/PycharmProjects/FPN_TF/tools/inference_image/2.jpg")
 126     img_array = np.array(img_array, np.float32) - np.array(cfgs.PIXEL_MEAN)
 127     boxes = np.array(
 128         [[200, 200, 500, 500],
 129          [300, 300, 400, 400],
 130          [200, 200, 400, 400]]
 131     )
 132
 133     # test only draw boxes
 134     labes = np.ones(shape=[len(boxes), ], dtype=np.float32) * ONLY_DRAW_BOXES
 135     scores = np.zeros_like(labes)
 136     imm = draw_boxes_with_label_and_scores(img_array, boxes, labes ,scores)
 137     # imm = np.array(imm)
 138
 139     cv2.imshow("te", imm)
 140
 141     # test only draw scores
 142     labes = np.ones(shape=[len(boxes), ], dtype=np.float32) * ONLY_DRAW_BOXES_WITH_SCORES
 143     scores = np.random.rand((len(boxes))) * 10
 144     imm2 = draw_boxes_with_label_and_scores(img_array, boxes, labes, scores)
 145
 146     cv2.imshow("te2", imm2)
 147     # test draw label and scores
 148
 149     labels = np.arange(1, 4)
 150     imm3 = draw_boxes_with_label_and_scores(img_array, boxes, labels, scores)
 151     cv2.imshow("te3", imm3)
 152
 153     cv2.waitKey(0)
 154
 155
 156
 157
 158
 159
 160