物體檢測標(biāo)注的小工具 - 在Python下使用OpenCV的技巧教程與典型應(yīng)用案例

6.4.3 代碼：物體檢測標(biāo)注的小工具

基于上面兩小節(jié)的基本使用，就能和OpenCV的基本繪圖功能就能實(shí)現(xiàn)一個超級簡單的物體框標(biāo)注小工具了?；舅悸肥菍σ獦?biāo)注的圖像建立一個窗口循環(huán)，然后每次循環(huán)的時(shí)候?qū)D像進(jìn)行一次拷貝。鼠標(biāo)在畫面上畫框的操作，以及已經(jīng)畫好的框的相關(guān)信息在全局變量中保存，并且在每個循環(huán)中根據(jù)這些信息，在拷貝的圖像上再畫一遍，然后顯示這份拷貝的圖像。

基于這種實(shí)現(xiàn)思路，使用上我們采用一個盡量簡化的設(shè)計(jì)：

- 輸入是一個文件夾，下面包含了所有要標(biāo)注物體框的圖片。如果圖片中標(biāo)注了物體，則生成一個相同名稱加額外后綴名的文件保存標(biāo)注信息。

- 標(biāo)注的方式是按下鼠標(biāo)左鍵選擇物體框的左上角，松開鼠標(biāo)左鍵選擇物體框的右下角，鼠標(biāo)右鍵刪除上一個標(biāo)注好的物體框。所有待標(biāo)注物體的類別，和標(biāo)注框顏色由用戶自定義，如果沒有定義則默認(rèn)只標(biāo)注一種物體，定義該物體名稱叫“Object”。

- 方向鍵的←和→用來遍歷圖片，↑和↓用來選擇當(dāng)前要標(biāo)注的物體，Delete鍵刪除一張圖片和對應(yīng)的標(biāo)注信息。

每張圖片的標(biāo)注信息，以及自定義標(biāo)注物體和顏色的信息，用一個元組表示，第一個元素是物體名字，第二個元素是代表BGR顏色的tuple或者是代表標(biāo)注框坐標(biāo)的元組。對于這種并不復(fù)雜復(fù)雜的數(shù)據(jù)結(jié)構(gòu)，我們直接利用Python的repr()函數(shù)，把數(shù)據(jù)結(jié)構(gòu)保存成機(jī)器可讀的字符串放到文件里，讀取的時(shí)候用eval()函數(shù)就能直接獲得數(shù)據(jù)。這樣的方便之處在于不需要單獨(dú)寫個格式解析器。如果需要可以在此基礎(chǔ)上再編寫一個轉(zhuǎn)換工具就能夠轉(zhuǎn)換成常見的Pascal VOC的標(biāo)注格式或是其他的自定義格式。

在這些思路和設(shè)計(jì)下，我們定義標(biāo)注信息文件的格式的例子如下：
('Hill', ((221, 163), (741, 291)))
('Horse', ((465, 430), (613, 570)))

元組中第一項(xiàng)是物體名稱，第二項(xiàng)是標(biāo)注框左上角和右下角的坐標(biāo)。這里之所以不把標(biāo)注信息的數(shù)據(jù)直接用pickle保存，是因?yàn)閿?shù)據(jù)本身不會很復(fù)雜，直接保存還有更好的可讀性。自定義標(biāo)注物體和對應(yīng)標(biāo)注框顏色的格式也類似，不過更簡單些，因?yàn)槔ㄌ柨梢圆粚懀唧w如下：

'Horse', (255, 255, 0)
'Hill', (0, 255, 255)
'DiaoSi', (0, 0, 255)

第一項(xiàng)是物體名稱，第二項(xiàng)是物體框的顏色。使用的時(shí)候把自己定義好的內(nèi)容放到一個文本里，然后保存成和待標(biāo)注文件夾同名，后綴名為labels的文件。比如我們在一個叫samples的文件夾下放上一些草原的照片，然后自定義一個samples.labels的文本文件。把上段代碼的內(nèi)容放進(jìn)去，就定義了小山頭的框?yàn)辄S色，駿馬的框?yàn)榍嗌?，以及紅色的屌絲。基于以上，標(biāo)注小工具的代碼如下：
import os
import cv2

# tkinter是Python內(nèi)置的簡單GUI庫，實(shí)現(xiàn)一些比如打開文件夾，確認(rèn)刪除等操作十分方便
from tkFileDialog import askdirectory
from tkMessageBox import askyesno

# 定義標(biāo)注窗口的默認(rèn)名稱
WINDOW_NAME = 'Simple Bounding Box Labeling Tool'

# 定義畫面刷新的大概幀率（是否能達(dá)到取決于電腦性能）
FPS = 24

# 定義支持的圖像格式
SUPPOTED_FORMATS = ['jpg', 'jpeg', 'png']

# 定義默認(rèn)物體框的名字為Object，顏色藍(lán)色，當(dāng)沒有用戶自定義物體時(shí)用默認(rèn)物體
DEFAULT_COLOR = {'Object': (255, 0, 0)}

# 定義灰色，用于信息顯示的背景和未定義物體框的顯示
COLOR_GRAY = (192, 192, 192)

# 在圖像下方多出BAR_HEIGHT這么多像素的區(qū)域用于顯示文件名和當(dāng)前標(biāo)注物體等信息
BAR_HEIGHT = 16