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')把座標軸關掉😯