# 上下左右,ESC及刪除鍵對(duì)應(yīng)的cv.waitKey()的返回值
# 注意這個(gè)值根據(jù)操作系統(tǒng)不同有不同,可以通過6.4.2中的代碼獲取
KEY_UP = 65362
KEY_DOWN = 65364
KEY_LEFT = 65361
KEY_RIGHT = 65363
KEY_ESC = 27
KEY_DELETE = 65535
# 空鍵用于默認(rèn)循環(huán)
KEY_EMPTY = 0
get_bbox_name = '{}.bbox'.format
# 定義物體框標(biāo)注工具類
class SimpleBBoxLabeling:
def __init__(self, data_dir, fps=FPS, window_name=None):
self._data_dir = data_dir
self.fps = fps
self.window_name = window_name if window_name else WINDOW_NAME
#pt0是正在畫的左上角坐標(biāo),pt1是鼠標(biāo)所在坐標(biāo)
self._pt0 = None
self._pt1 = None
# 表明當(dāng)前是否正在畫框的狀態(tài)標(biāo)記
self._drawing = False
# 當(dāng)前標(biāo)注物體的名稱
self._cur_label = None
# 當(dāng)前圖像對(duì)應(yīng)的所有已標(biāo)注框
self._bboxes = []
# 如果有用戶自定義的標(biāo)注信息則讀取,否則用默認(rèn)的物體和顏色
label_path = '{}.labels'.format(self._data_dir)
self.label_colors = DEFAULT_COLOR if not os.path.exists(label_path) else self.load_labels(label_path)
# 獲取已經(jīng)標(biāo)注的文件列表和還未標(biāo)注的文件列表
imagefiles = [x for x in os.listdir(self._data_dir) if x[x.rfind('.') + 1:].lower() in SUPPOTED_FORMATS]
labeled = [x for x in imagefiles if os.path.exists(get_bbox_name(x))]
to_be_labeled = [x for x in imagefiles if x not in labeled]
# 每次打開一個(gè)文件夾,都自動(dòng)從還未標(biāo)注的第一張開始
self._filelist = labeled + to_be_labeled
self._index = len(labeled)
if self._index > len(self._filelist) - 1:
self._index = len(self._filelist) - 1
# 鼠標(biāo)回調(diào)函數(shù)
def _mouse_ops(self, event, x, y, flags, param):
# 按下左鍵時(shí),坐標(biāo)為左上角,同時(shí)表明開始畫框,改變drawing標(biāo)記為True
if event == cv2.EVENT_LBUTTONDOWN:
self._drawing = True
self._pt0 = (x, y)
# 左鍵抬起,表明當(dāng)前框畫完了,坐標(biāo)記為右下角,并保存,同時(shí)改變drawing標(biāo)記為False
elif event == cv2.EVENT_LBUTTONUP:
self._drawing = False
self._pt1 = (x, y)
self._bboxes.append((self._cur_label, (self._pt0, self._pt1)))
# 實(shí)時(shí)更新右下角坐標(biāo)方便畫框
elif event == cv2.EVENT_MOUSEMOVE:
self._pt1 = (x, y)
# 鼠標(biāo)右鍵刪除最近畫好的框
elif event == cv2.EVENT_RBUTTONUP:
if self._bboxes:
self._bboxes.pop()
# 清除所有標(biāo)注框和當(dāng)前狀態(tài)
def _clean_bbox(self):
self._pt0 = None
self._pt1 = None
self._drawing = False
self._bboxes = []
# 畫標(biāo)注框和當(dāng)前信息的函數(shù)
def _draw_bbox(self, img):
# 在圖像下方多出BAR_HEIGHT這么多像素的區(qū)域用于顯示文件名和當(dāng)前標(biāo)注物體等信息
h, w = img.shape[:2]
canvas = cv2.copyMakeBorder(img, 0, BAR_HEIGHT, 0, 0, cv2.BORDER_CONSTANT, value=COLOR_GRAY)
# 正在標(biāo)注的物體信息,如果鼠標(biāo)左鍵已經(jīng)按下,則顯示兩個(gè)點(diǎn)坐標(biāo),否則顯示當(dāng)前待標(biāo)注物體的名稱
label_msg = '{}: {}, {}'.format(self._cur_label, self._pt0, self._pt1)
if self._drawing
else 'Current label: {}'.format(self._cur_label)
# 顯示當(dāng)前文件名,文件個(gè)數(shù)信息
msg = '{}/{}: {} | {}'.format(self._index + 1, len(self._filelist), self._filelist[self._index], label_msg)
cv2.putText(canvas, msg, (1, h+12),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 0, 0), 1)
# 畫出已經(jīng)標(biāo)好的框和對(duì)應(yīng)名字
for label, (bpt0, bpt1) in self._bboxes:
label_color = self.label_colors[label] if label in self.label_colors else COLOR_GRAY
cv2.rectangle(canvas, bpt0, bpt1, label_color, thickness=2)
cv2.putText(canvas, label, (bpt0[0]+3, bpt0[1]+15),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, label_color, 2)
# 畫正在標(biāo)注的框和對(duì)應(yīng)名字
if self._drawing:
label_color = self.label_colors[self._cur_label] if self._cur_label in self.label_colors else COLOR_GRAY
if self._pt1[0] >= self._pt0[0] and self._pt1[1] >= self._pt0[1]:
cv2.rectangle(canvas, self._pt0, self._pt1, label_color, thickness=2)
cv2.putText(canvas, self._cur_label, (self._pt0[0] + 3, self._pt0[1] + 15),
cv2.FONT_HERSHEY_SIMPLEX,
0.5, label_color, 2)
return canvas
# 利用repr()導(dǎo)出標(biāo)注框數(shù)據(jù)到文件
@staticmethod
def export_bbox(filepath, bboxes):
if bboxes:
with open(filepath, 'w') as f:
for bbox in bboxes:
line = repr(bbox) + '
'
f.write(line)
elif os.path.exists(filepath):
os.remove(filepath)
# 利用eval()讀取標(biāo)注框字符串到數(shù)據(jù)
@staticmethod
def load_bbox(filepath):
bboxes = []
with open(filepath, 'r') as f:
line = f.readline().rstrip()
while line:
bboxes.append(eval(line))
line = f.readline().rstrip()
return bboxes
評(píng)論
查看更多