X-Git-Url: https://gerrit.akraino.org/r/gitweb?a=blobdiff_plain;f=example-apps%2FPDD%2Fpcb-defect-detection%2Flibs%2Fbox_utils%2Fdraw_box_in_img.py;fp=example-apps%2FPDD%2Fpcb-defect-detection%2Flibs%2Fbox_utils%2Fdraw_box_in_img.py;h=0000000000000000000000000000000000000000;hb=3ed2c61d9d7e7916481650c41bfe5604f7db22e9;hp=c80c5d2af0a2d26560ebbd5026e4dc6b1d6019ed;hpb=e6d40ddb2640f434a9d7d7ed99566e5e8fa60cc1;p=ealt-edge.git diff --git a/example-apps/PDD/pcb-defect-detection/libs/box_utils/draw_box_in_img.py b/example-apps/PDD/pcb-defect-detection/libs/box_utils/draw_box_in_img.py deleted file mode 100755 index c80c5d2..0000000 --- a/example-apps/PDD/pcb-defect-detection/libs/box_utils/draw_box_in_img.py +++ /dev/null @@ -1,160 +0,0 @@ -# -*- coding: utf-8 -*- - -from __future__ import absolute_import, print_function, division - -import numpy as np - -from PIL import Image, ImageDraw, ImageFont -import cv2 - -from libs.configs import cfgs -from libs.label_name_dict.label_dict import LABEl_NAME_MAP -NOT_DRAW_BOXES = 0 -ONLY_DRAW_BOXES = -1 -ONLY_DRAW_BOXES_WITH_SCORES = -2 - -STANDARD_COLORS = [ - 'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque', - 'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite', - 'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan', - 'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange', - 'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet', - 'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite', - 'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod', - 'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki', - 'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue', - 'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey', - 'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue', - 'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime', - 'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid', - 'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen', - 'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin', - 'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed', - 'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed', - 'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple', - 'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown', - 'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue', - 'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow', - 'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White', - 'WhiteSmoke', 'Yellow', 'YellowGreen', 'LightBlue', 'LightGreen' -] -FONT = ImageFont.load_default() - - -def draw_a_rectangel_in_img(draw_obj, box, color, width): - ''' - use draw lines to draw rectangle. since the draw_rectangle func can not modify the width of rectangle - :param draw_obj: - :param box: [x1, y1, x2, y2] - :return: - ''' - x1, y1, x2, y2 = box[0], box[1], box[2], box[3] - top_left, top_right = (x1, y1), (x2, y1) - bottom_left, bottom_right = (x1, y2), (x2, y2) - - draw_obj.line(xy=[top_left, top_right], - fill=color, - width=width) - draw_obj.line(xy=[top_left, bottom_left], - fill=color, - width=width) - draw_obj.line(xy=[bottom_left, bottom_right], - fill=color, - width=width) - draw_obj.line(xy=[top_right, bottom_right], - fill=color, - width=width) - - -def only_draw_scores(draw_obj, box, score, color): - - x, y = box[0], box[1] - draw_obj.rectangle(xy=[x, y-10, x+100, y], - fill=color) - draw_obj.text(xy=(x, y-10), - text="obj:" +str(round(score, 2)), - fill='white', - font=FONT) - - -def draw_label_with_scores(draw_obj, box, label, score, color): - x, y = box[0], box[1] - draw_obj.rectangle(xy=[x, y-10, x+100, y], - fill=color) - - txt = LABEl_NAME_MAP[label] + ':' + str(round(score, 2)) - draw_obj.text(xy=(x, y-10), - text=txt, - fill='white', - font=FONT) - - -def draw_boxes_with_label_and_scores(img_array, boxes, labels, scores): - - img_array = img_array + np.array(cfgs.PIXEL_MEAN) - img_array.astype(np.float32) - boxes = boxes.astype(np.int64) - labels = labels.astype(np.int32) - img_array = np.array(img_array * 255 / np.max(img_array), dtype=np.uint8) - - img_obj = Image.fromarray(img_array) - raw_img_obj = img_obj.copy() - - draw_obj = ImageDraw.Draw(img_obj) - num_of_objs = 0 - for box, a_label, a_score in zip(boxes, labels, scores): - - if a_label != NOT_DRAW_BOXES: - num_of_objs += 1 - #draw_a_rectangel_in_img(draw_obj, box, color=STANDARD_COLORS[a_label], width=3) - draw_a_rectangel_in_img(draw_obj, box, color='Red', width=3) - if a_label == ONLY_DRAW_BOXES: # -1 - continue - elif a_label == ONLY_DRAW_BOXES_WITH_SCORES: # -2 - only_draw_scores(draw_obj, box, a_score, color='Yellow') - continue - else: - draw_label_with_scores(draw_obj, box, a_label, a_score, color='DeepSkyBlue') - - out_img_obj = Image.blend(raw_img_obj, img_obj, alpha=0.7) - - return np.array(out_img_obj) - - -if __name__ == '__main__': - img_array = cv2.imread("/home/yjr/PycharmProjects/FPN_TF/tools/inference_image/2.jpg") - img_array = np.array(img_array, np.float32) - np.array(cfgs.PIXEL_MEAN) - boxes = np.array( - [[200, 200, 500, 500], - [300, 300, 400, 400], - [200, 200, 400, 400]] - ) - - # test only draw boxes - labes = np.ones(shape=[len(boxes), ], dtype=np.float32) * ONLY_DRAW_BOXES - scores = np.zeros_like(labes) - imm = draw_boxes_with_label_and_scores(img_array, boxes, labes ,scores) - # imm = np.array(imm) - - cv2.imshow("te", imm) - - # test only draw scores - labes = np.ones(shape=[len(boxes), ], dtype=np.float32) * ONLY_DRAW_BOXES_WITH_SCORES - scores = np.random.rand((len(boxes))) * 10 - imm2 = draw_boxes_with_label_and_scores(img_array, boxes, labes, scores) - - cv2.imshow("te2", imm2) - # test draw label and scores - - labels = np.arange(1, 4) - imm3 = draw_boxes_with_label_and_scores(img_array, boxes, labels, scores) - cv2.imshow("te3", imm3) - - cv2.waitKey(0) - - - - - - -