Python可視化目標檢測框的實現代碼

1 引言

隨著計算機視覺算法工程師的內卷,從事目標檢測的小夥伴們越來越多瞭.
很多時候我們費瞭九牛二虎之力訓練瞭一版模型,可是可視化出來的效果平淡無奇.
是不是有點太不給力啦,作為計算機視覺工程師,我們是不是應該關註下如何優雅地可視化我們模型地檢測結果呢?

2 舉個栗子

最常用的可視化目標檢測結果的就是我們所說的矩形框,矩形框的畫法也可以分為好多中,我們以下圖進行說明:

請添加圖片描述

我們以上圖皮卡丘的矩形檢測框為例,左上為我們常用的不帶標簽的輸出框,右上為YOLO系列美化帶標簽後的矩形框,左下為四個角點美化後的矩形框,右下為帶標簽的角點美化矩形框.

3 實現

3.1 函數講解

在opencv中,我們通常使用cv2.rectangle函數進行矩形框的繪制,該函數的一般形式如下:

cv2.rectangle(image, start_point, end_point, color, thickness)

相應的參數含義如下:

  • 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)
    cv2.line(draw_img, (bbox[0], bbox[3]), (bbox[0], bbox[3] - length), corner_color, thickness=3)
    # Bottom Right
    cv2.line(draw_img, (bbox[2], bbox[3]), (bbox[2] - length, bbox[3]), corner_color, thickness=3)
    cv2.line(draw_img, (bbox[2], bbox[3]), (bbox[2], bbox[3] - length), corner_color, thickness=3)

上述函數參數解釋如下:

  • draw_img 輸入圖像
  • bbox 目標檢測框 形式(x1,y1,x2,y2)
  • length 直線長度
  • corner_color 直線顏色

運行結果如下:

請添加圖片描述

左側為原圖,右側為添加角點美化後的結果圖.

3.5 綜合效果

綜合上述標簽美化和角點美化,可以得到二者綜合後的效果圖如下:

請添加圖片描述

4 透明效果實現

通過cv2.rectangle 函數設置 thickness 我們可以得到填充後的圖像,進而利用opencv中的 cv2.add_weight函數可以實現透明效果.
示例如下:

請添加圖片描述

左側為img , 右側為填充後的draw_img,那麼利用add_weight 進行加權的代碼如下:

alpha = 0.8
gamma = 0
out_img = cv2.addWeighted(img,alpha,draw_img,1-alpha,gamma)

運行後的效果圖如下:

請添加圖片描述

左側為 alpha=0.5的效果圖,右側為alpha=0.8的效果圖.
和原圖放到一起進行對比,如下:

請添加圖片描述

是不是看上去優雅多瞭…

5 擴展應用

利用上述原理,我們可以方便地將透明效果從矩形框遷移至不規則封閉區域, 樣例如下:

請添加圖片描述

左側為我們輸入分割網絡的原始圖像,右側為我們分割網絡可行駛區域的推理圖,我們使用alpha=0.6對上述兩幅圖進行加權,得到結果如下:

請添加圖片描述

6 總結

本文實現瞭目標檢測矩形框的標簽優化以及角點優化,並給出瞭透明化的具體原理和實際擴展應用,以及詳細的代碼實現.

您學廢瞭嗎?

7 參考

鏈接一

鏈接二

到此這篇關於Python如何優雅地可視化目標檢測框的文章就介紹到這瞭,更多相關Python目標檢測框內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet!

推薦閱讀: