import cv2 import webcolors import os import uuid import numpy as np STANDARD_COLORS = [ 'LawnGreen', 'Chartreuse', 'Aqua', 'Beige', 'Azure', 'BlanchedAlmond', 'Bisque', 'Aquamarine', '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', 'AliceBlue', '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' ] def from_colorname_to_bgr(color): rgb_color = webcolors.name_to_rgb(color) result = (rgb_color.blue, rgb_color.green, rgb_color.red) return result def standard_to_bgr(list_color_name): standard = [] for i in range(len(list_color_name) - 36): # -36 used to match the len(obj_list) standard.append(from_colorname_to_bgr(list_color_name[i])) return standard def get_index_label(label, obj_list): index = int(obj_list.index(label)) return index def plot_one_box(img, coord, label=None, score=None, color=None, line_thickness=None): tl = line_thickness or int(round(0.001 * max(img.shape[0:2]))) # line thickness color = color c1, c2 = (int(coord[0]), int(coord[1])), (int(coord[2]), int(coord[3])) cv2.rectangle(img, c1, c2, color, thickness=tl) if label: tf = max(tl - 2, 1) # font thickness s_size = cv2.getTextSize(str('{:.0%}'.format(score)), 0, fontScale=float(tl) / 3, thickness=tf)[0] t_size = cv2.getTextSize(label, 0, fontScale=float(tl) / 3, thickness=tf)[0] c2 = c1[0] + t_size[0] + s_size[0] + 15, c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1) # filled cv2.putText(img, '{}: {:.0%}'.format(label, score), (c1[0], c1[1] - 2), 0, float(tl) / 3, [0, 0, 0], thickness=tf, lineType=cv2.FONT_HERSHEY_SIMPLEX) color_list = standard_to_bgr(STANDARD_COLORS) def display(preds, imgs, obj_list, imshow=True, imwrite=False): for i in range(len(imgs)): if len(preds[i]['rois']) == 0: continue imgs[i] = imgs[i].copy() for j in range(len(preds[i]['rois'])): (x1, y1, x2, y2) = preds[i]['rois'][j].astype(np.int) obj = obj_list[preds[i]['class_ids'][j]] score = float(preds[i]['scores'][j]) plot_one_box(imgs[i], [x1, y1, x2, y2], label=obj, score=score, color=color_list[get_index_label(obj, obj_list)]) if imshow: cv2.imshow('img', imgs[i]) cv2.waitKey(0) if imwrite: os.makedirs('test/', exist_ok=True) cv2.imwrite(f'test/{uuid.uuid4().hex}.jpg', imgs[i])