Python如何优雅地可视化目标检测框

卓不凡

1 引言

随着计算机视觉算法工程师的内卷,从事目标检测的小伙伴们越来越多了.
很多时候我们费了九牛二虎之力训练了一版模型,可是可视化出来的效果平淡无奇.
是不是有点太不给力啦,作为计算机视觉工程师,我们是不是应该关注下如何优雅地可视化我们模型地检测结果呢?

2 举个栗子

最常用的可视化目标检测结果的就是我们所说的矩形框,矩形框的画法也可以分为好多中,我们以下图进行说明:

上图皮卡丘的矩形检测效果示例中:
左上为我们常用的不带标签的输出框,右上为YOLO系列美化带标签后的矩形框.
左下为四个角点美化后的矩形框,右下为带标签的角点美化矩形框.

3 实现

3.1 函数讲解

在opencv中,我们通常使用cv2.rectangle函数进行矩形框的绘制,该函数的一般形式如下:

相应的参数含义如下:

image: 输入图像
start_point: 矩形框左上点坐标
end_point: 矩形框右下点坐标
color: 矩形框颜色默认BGR顺序
thickness: 线的粗细,其中 -1 代表填充整个矩形

3.2 读入图像

我们使用上述函数,来画我们的示例图像:

img_name = './pikachu.jpg'
img = cv2.imread(img_name)
box = [ 140, 16,468,390, "pikachu"]
box_color = (255,0,255)   
cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color=box_color, thickness=2)

结果如下:

左侧为我们的原图,右侧为我们画框的效果图.

3.3 标签美化

接下来我们来给矩形框添加标签,我们观察上述画图函数,注意最后一个参数 thickness ,
如果此值等于-1,那么将对矩形框执行填充效果.
基于此,我们来写标签美化代码,如下:

def draw_label_type(draw_img,bbox,label_color):
    label = str(bbox[-1])
    labelSize = cv2.getTextSize(label + '0', cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
    if bbox[1] - labelSize[1] - 3 < 0:
        cv2.rectangle(draw_img,
                      (bbox[0], bbox[1] + 2),
                      (bbox[0] + labelSize[0], bbox[1] + labelSize[1] + 3),
                      color=label_color,
                      thickness=-1
        cv2.putText(draw_img, label,
                    (bbox[0], bbox[1] + labelSize + 3),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 0, 0),
                    thickness=1
    else:
        cv2.rectangle(draw_img,
                      (bbox[0], bbox[1] - labelSize[1] - 3),
                      (bbox[0] + labelSize[0], bbox[1] - 3),
                      color=label_color,
                      thickness=-1
        cv2.putText(draw_img, label,
                    (bbox[0], bbox[1] - 3),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 0, 0),
                    thickness=1
                    )

上述代码中, 首先计算标签文本的大小, 然后据标签文本大小进行矩形填充, 最后使用函数 cv2.putText 来画出对应的标签文本.
运行效果如下:

上图中左侧为原图,右侧为添加文本标签后的结果图.

3.4 角点美化

上述添加完标签后,由于标签框和目标矩形框颜色一致,边界处不太容易区分,这里添加对角点美化的代码,代码如下:

def draw_box_corner(draw_img,bbox,length,corner_color):
    # Top Left
    cv2.line(draw_img, (bbox[0], bbox[1]), (bbox[0] + length, bbox[1]), corner_color, thickness=3)
    cv2.line(draw_img, (bbox[0], bbox[1]), (bbox[0], bbox[1] + length), corner_color, thickness=3)
    # Top Right
    cv2.line(draw_img, (bbox[2], bbox[1]), (bbox[2] - length, bbox[1]), corner_color, thickness=3)
    cv2.line(draw_img, (bbox[2], bbox[1]), (bbox[2], bbox[1] + length), corner_color, thickness=3)
    # Bottom Left
    cv2.line(draw_img, (bbox[0], bbox[3]), (bbox[0] + length, bbox[3]), corner_color, thickness=3)