如有错误,恳请指出。
在之前的内容中,基本已经把trick与处理流程讲了一遍,这两个函数只是进行缩放与画框,也谈不上是一个技巧,也可以说是一个工具。
所以下面介绍两个函数:
- sacle_coords函数:将当前的label坐标位置尺寸还原为原图的坐标位置尺寸
- draw_box函数:将当前的预测边界框显示在原图上
- 1. sacle_coords函数
- 2. draw_box函数
- 主要思路:
获取缩放后的图像与缩放前图像的一个比例,根据比例计算出上下左右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
如果实在看不懂也没有关系,根据参数传入相对于的数据就可以直接使用了。需要传入的数据有图像本身,边界框信息,类别信息,置信度信息,字典对信息(类别索引与类别名称的键值对)。剩下的阈值与框粗细自行选择设置。
所以,总的来说,以上两个函数也可以看成是一个工具来使用。
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)