目标检测的Tricks | 【Trick11】label的缩放与显示

目标检测的Tricks | 【Trick11】label的缩放与显示,第1张


如有错误,恳请指出。


在之前的内容中,基本已经把trick与处理流程讲了一遍,这两个函数只是进行缩放与画框,也谈不上是一个技巧,也可以说是一个工具。

所以下面介绍两个函数:

  • sacle_coords函数:将当前的label坐标位置尺寸还原为原图的坐标位置尺寸
  • draw_box函数:将当前的预测边界框显示在原图上

文章目录
  • 1. sacle_coords函数
  • 2. draw_box函数

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

如果实在看不懂也没有关系,根据参数传入相对于的数据就可以直接使用了。需要传入的数据有图像本身,边界框信息,类别信息,置信度信息,字典对信息(类别索引与类别名称的键值对)。剩下的阈值与框粗细自行选择设置。

所以,总的来说,以上两个函数也可以看成是一个工具来使用。

欢迎分享,转载请注明来源:内存溢出

原文地址: https://outofmemory.cn/langs/738487.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-04-28
下一篇 2022-04-28

发表评论

登录后才能评论

评论列表(0条)

保存