# 上下左右，ESC及刪除鍵對應的cv.waitKey()的返回值
# 注意這個值根據操作系統不同有不同，可以通過6.4.2中的代碼獲取
KEY_UP = 65362
KEY_DOWN = 65364
KEY_LEFT = 65361
KEY_RIGHT = 65363
KEY_ESC = 27
KEY_DELETE = 65535

# 空鍵用于默認循環
KEY_EMPTY = 0

get_bbox_name = '{}.bbox'.format

# 定義物體框標注工具類
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是正在畫的左上角坐標，pt1是鼠標所在坐標
self._pt0 = None
self._pt1 = None

# 表明當前是否正在畫框的狀態標記
self._drawing = False

# 當前標注物體的名稱
self._cur_label = None

# 當前圖像對應的所有已標注框
self._bboxes = []

# 如果有用戶自定義的標注信息則讀取，否則用默認的物體和顏色
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)

# 獲取已經標注的文件列表和還未標注的文件列表
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]

# 每次打開一個文件夾，都自動從還未標注的第一張開始
self._filelist = labeled + to_be_labeled
self._index = len(labeled)
if self._index > len(self._filelist) - 1:
self._index = len(self._filelist) - 1

# 鼠標回調函數
def _mouse_ops(self, event, x, y, flags, param):

# 按下左鍵時，坐標為左上角，同時表明開始畫框，改變drawing標記為True
if event == cv2.EVENT_LBUTTONDOWN:
self._drawing = True
self._pt0 = (x, y)

# 左鍵抬起，表明當前框畫完了，坐標記為右下角，并保存，同時改變drawing標記為False
elif event == cv2.EVENT_LBUTTONUP:
self._drawing = False
self._pt1 = (x, y)
self._bboxes.append((self._cur_label, (self._pt0, self._pt1)))

# 實時更新右下角坐標方便畫框
elif event == cv2.EVENT_MOUSEMOVE:
self._pt1 = (x, y)

# 鼠標右鍵刪除最近畫好的框
elif event == cv2.EVENT_RBUTTONUP:
if self._bboxes:
self._bboxes.pop()

# 清除所有標注框和當前狀態
def _clean_bbox(self):
self._pt0 = None
self._pt1 = None
self._drawing = False
self._bboxes = []

# 畫標注框和當前信息的函數
def _draw_bbox(self, img):

# 在圖像下方多出BAR_HEIGHT這么多像素的區域用于顯示文件名和當前標注物體等信息
h, w = img.shape[:2]
canvas = cv2.copyMakeBorder(img, 0, BAR_HEIGHT, 0, 0, cv2.BORDER_CONSTANT, value=COLOR_GRAY)

# 正在標注的物體信息，如果鼠標左鍵已經按下，則顯示兩個點坐標，否則顯示當前待標注物體的名稱
label_msg = '{}: {}, {}'.format(self._cur_label, self._pt0, self._pt1)
if self._drawing
else 'Current label: {}'.format(self._cur_label)

# 顯示當前文件名，文件個數信息
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)

# 畫出已經標好的框和對應名字
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)

# 畫正在標注的框和對應名字
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()導出標注框數據到文件
@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()讀取標注框字符串到數據
@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