1. sacle_coords函数
主要思路：
获取缩放后的图像与缩放前图像的一个比例，根据比例计算出上下左右pad的数目。然后对边界框减去pad的数目，再根据比例缩放，就获得了一个缩放后的边界框。最后进行一个边界截断，以免溢出。返回缩放处理好的的预测框。
参考代码
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): """ 将预测的坐标信息转换回原图尺度 :param img1_shape: 缩放后的图像尺度 :param coords: 预测的box信息 :param img0_shape: 缩放前的图像尺度 :param ratio_pad: 缩放过程中的缩放比例以及pad :return: """ # Rescale coords (xyxy) from img1_shape to img0_shape if ratio_pad is None: # calculate from img0_shape gain = max(img1_shape) / max(img0_shape) # gain = old / new pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] coords[:, [0, 2]] -= pad[0] # x padding coords[:, [1, 3]] -= pad[1] # y padding coords[:, :4] /= gain clip_coords(coords, img0_shape) return coords def clip_coords(boxes, img_shape): # Clip bounding xyxy bounding boxes to image shape (height, width) boxes[:, 0].clamp_(0, img_shape[1]) # x1 boxes[:, 1].clamp_(0, img_shape[0]) # y1 boxes[:, 2].clamp_(0, img_shape[1]) # x2 boxes[:, 3].clamp_(0, img_shape[0]) # y2
其中的重点是根据比例来推断出上下左右的pad数目，边界框的数值需要减去这个pad偏移，才可以对应真正的原图上的大小。
2. draw_box函数
如图所示，对于这些检测出来的边界框，其实就是需要根据预测好的预测框坐标然后在原图上画出来，才会得到这些一个个的框，每个类别对于一种颜色，其中还需要显示类别名称与置信度的大小。
这里还会进行最后一轮的过滤，当nms处理后的预测框，当置信度还存在低于阈值0.1的预测框，这里会直接忽视不显示。然后有个小方法，对于每个预测框构建两个字典，一个是用来存储要显示的字符串，一个是用来存储要显示的框颜色（框颜色是根据类别来选择的）。如下所示：
之后需要的做的事情，就是遍历字典中的每一个预测框，根据坐标信息与文本信息（类别名称：置信度）来在相应的位置绘制出来。重点函数如下所示：
# 逐个在原图上添加预测框信息 for box, color in box_to_color_map.items(): # 两个点可以确定一个框 xmin, ymin, xmax, ymax = box (left, right, top, bottom) = (xmin * 1, xmax * 1, ymin * 1, ymax * 1) # 这里是划线，由5个点画4条线来确定一个框，输入的是一个列表 # width: 划线的宽粗控制 # fill: 划线的颜色填充 draw.line([(left, top), (left, bottom), (right, bottom), (right, top), (left, top)], width=line_thickness, fill=color) # 文本信息 draw_text(draw, box_to_display_str_map, box, left, right, top, bottom, color)
整个draw_box代码如下所示，我做好了注释的：
import collections from PIL import Image import PIL.ImageDraw as ImageDraw import PIL.ImageFont as ImageFont import numpy as np 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' ] # 过滤低于阈值(thresh=0.1)的预测框 def filter_low_thresh(boxes, scores, classes, category_index, thresh, box_to_display_str_map, box_to_color_map): """ 1、过滤掉scores低于thresh的anchor; 2、为每个anchor生成显示信息和框框颜色并分别保存在box_to_display_str_map和box_to_color_map中 :param boxes: 最终预测结果 (anchor_nums, x1+y1+x2+y2)=(7, 4) (相对原图的预测结果) 分类别且按score从大到小排列 :param scores: 所有预测anchors的得分 (7) 分类别且按score从大到小排列 :param classes: 所有预测anchors的类别 (7) 分类别且按score从大到小排列 :param category_index: 所有类别的信息（从data/pascal_voc_classes.json中读出） :param thresh: 设置阈值（默认0.1），过滤掉score太低的anchor :param box_to_display_str_map: 拿来存放每个anchor的显示信息（list） 每个anchor: tuple(box) = list[显示信息] :param box_to_color_map: 拿来存放每个anchor的框框颜色 """ for i in range(boxes.shape[0]): # 如果低于阈值直接忽略 if scores[i] > thresh: # 转换成tuple类型: numpy -> list -> tuple box = tuple(boxes[i].tolist()) # 根据映射字典获取类别名称，如果没有指定赋值为'N/A' if classes[i] in category_index.keys(): class_name = category_index[classes[i]] else: class_name = 'N/A' # 显示的内容格式为: classname: score display_str = str(class_name) display_str = '{}: {}%'.format(display_str, int(100 * scores[i])) # 键值为预测框: 分别对应着需要显示的内容与边界框的颜色 box_to_display_str_map[box].append(display_str) box_to_color_map[box] = STANDARD_COLORS[ classes[i] % len(STANDARD_COLORS)] # 网络输出概率已经排序过，当遇到一个不满足后面的肯定不满足 else: break def draw_text(draw, box_to_display_str_map, box, left, right, top, bottom, color): """ :param draw: 一个可以在给定图像(image)上绘图的对象 :param box_to_display_str_map: 每个anchor的显示信息 :param box: 当前anchor的预测信息 (xyxy) :param left: anchor的left :param right: anchor的right :param top: anchor的top :param bottom: anchor的bottom :param color: 当前anchor的信息颜色/anchor框框颜色 :return: """ # 字体选择 try: font = ImageFont.truetype('arial.ttf', 20) except IOError: font = ImageFont.load_default() # 如果显示字符串的总高度添加到边界的顶部框超出图像顶部，将字符串堆叠在边界框下方而不是上面 display_str_heights = [font.getsize(ds)[1] for ds in box_to_display_str_map[box]] # 每个 display_str 的顶部和底部边距为 0.05x total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights) if top > total_display_str_height: text_bottom = top else: text_bottom = bottom + total_display_str_height # 反转列表并从下到上打印 for display_str in box_to_display_str_map[box][::-1]: # font.getsize函数可以获取当前文本信息的长度与高度(像素值) text_width, text_height = font.getsize(display_str) margin = np.ceil(0.05 * text_height) # 绘制文本信息后的背景 draw.rectangle([(left, text_bottom - text_height - 2 * margin), # [(331.0, 107.0), (403.0, 120.0)] (left + text_width, text_bottom)], fill=color) # 绘制文本信息 draw.text((left + margin, text_bottom - text_height - margin), # (332.0, 108.0) display_str, fill='black', font=font) text_bottom -= text_height - 2 * margin def draw_box(image, boxes, classes, scores, category_index, thresh=0.1, line_thickness=3): """ :param image: 原图 RGB (375, 500, 3) HWC numpy格式(array) img_o[:, :, ::-1]:BGR=>RGB :param boxes: 最终预测结果 (anchor_nums, x1+y1+x2+y2)=(7, 4) (相对原图的预测结果) 按score从大到小排列 numpy格式(array) :param classes: 所有预测anchors的类别 (7) 分类别且按score从大到小排列 numpy格式(array) :param scores: 所有预测anchors的得分 (7) 分类别且按score从大到小排列 numpy格式(array) :param category_index: 所有类别的信息（从data/pascal_voc_classes.json中读出） :param thresh: 设置阈值（默认0.1），过滤掉score太低的anchor :param line_thickness: 框框直线厚度 :return: """ # 根据边界框构建两个字典 box_to_display_str_map = collections.defaultdict(list) box_to_color_map = collections.defaultdict(str) # 过滤低阈值预测框 filter_low_thresh(boxes, scores, classes, category_index, thresh, box_to_display_str_map, box_to_color_map) # Draw all boxes onto image. if isinstance(image, np.ndarray): image = Image.fromarray(image) draw = ImageDraw.Draw(image) # im_width, im_height = image.size # 逐个在原图上添加预测框信息 for box, color in box_to_color_map.items(): # 两个点可以确定一个框 xmin, ymin, xmax, ymax = box (left, right, top, bottom) = (xmin * 1, xmax * 1, ymin * 1, ymax * 1) # 这里是划线，由5个点画4条线来确定一个框，输入的是一个列表 # width: 划线的宽粗控制 # fill: 划线的颜色填充 draw.line([(left, top), (left, bottom), (right, bottom), (right, top), (left, top)], width=line_thickness, fill=color) # 文本信息 draw_text(draw, box_to_display_str_map, box, left, right, top, bottom, color) return image
如果实在看不懂也没有关系，根据参数传入相对于的数据就可以直接使用了。需要传入的数据有图像本身，边界框信息，类别信息，置信度信息，字典对信息（类别索引与类别名称的键值对）。剩下的阈值与框粗细自行选择设置。
所以，总的来说，以上两个函数也可以看成是一个工具来使用。
