matplotlib在目标检测中绘制bbox的统一流程

嗯 感觉这块可能以后写检测算法都会用，就总结一下其中一些的API加深记忆吧。

目标检测中，我们一般是这样做的，我们定义的NMS方法一般是针对整个batch的，所以对于test集中每个batch的pred经过NMS后会得到这个batch最终的result了，一般是一个list 每个index对应一个(N,7)维的tensor，其中7就是4个坐标 置信度 最大class分数，最大class index（置信度可有可无，但一般都是这些信息）。然后我们要做的就是遍历batch的每一个img，根据每个img的这个N来绘制rectangle，最后保存图片。

预备知识

python有三种常用的图像库 opencv matplotlib PIL

plt.imread和PIL.Image.open读入的都是RGB顺序，而opencv中cv2.imread读入的是BGR通道顺序 ,所以如果要用plt显示的话需要通过cv2.cvtColor转换。其中PIL读入的是图片，如果要操作需要np.array() 其他两种得到的都是np数组。

绘制框

(1)为每个类别挑选一种独特的颜色 tab20b是一种20的color map

cmap = plt.get_camp('tab20b')
colors = [cmap(i) for i in np.linespace(0, 1, 20)]

np.linespace(start, end, step)

(2)对于batch中的每一张img，需要单独绘制一个figure

plt.figure()
fig, ax = plt.subplots(1)

plt.subplots()是一个函数，返回一个包含figure和axes对象的元组。因此，使用fig,ax = plt.subplots()将元组分解为fig和ax两个变量。如果为plt.subplots(1,2) ax可通过下表index来控制子图

(3)读入图片放入ax中

img = plt.imread(img_path)
ax.imshow(img)

注意opencv读入需要转换为rgb，PIL读入需要转为np ax.imshow是用来处理图片而不是show图片

(4)对于图片中的所有类别生成该图片的颜色

# detections list 每一个element为 (N, 7) tensor

for img_id, detection in enumerate(detections):
    unique_labels = detections[:, -1].cpu().unique()
    label_nums = len(unique_labels)
    bbox_color = random.sample(colors, label_nums)

(5）对检测结果中的每一个box，通过matplotlib.patches来绘制rectangle

for x1,y1, x2,y2,conf,cls_score, cls_index in detection:
    w = x2 - x1
    h = y2 - y1
    cls_color = bbox_color[int(np.where(unique_labels == int(cls_index))[0])]
    bbox = patches.Rectangle((x1,y1),w,h,linewidth=2,edge_color = color,facecolor='none')
    ax.add_patches(bbox)
    plt.text(
        x1,y1,class[int(cls_index)], color='white', verticalalignment ='top')

选择颜色，绘制矩形框用patches.Rectangle() 其中edge_color为框的边缘颜色，注意要把facecolor设置为none代表不填充。

plt.text在(x1,y1)处绘制text文本框，然后输入一个字符串代表text的字，color为text颜色，后面为对齐方式。

(6)绘制完所有bbox后保存图片

plt.savefig(filename, bbox_inches='tight', pad_inches= 0.0)

其中bbox_inches为tight代表figure自适应img size， pad_inches去掉空白，整体得到的图片不包含空白，同时别忘了通过plt.axis('off')把坐标轴关掉😯