天天看點

機器視覺 OpenCV—python 标注小工具(目标檢測)

一、擷取圖像并展示

# 擷取 picture_dir 下檔案的所有圖檔,并以24幀率顯示

import os
from itertools import cycle

filenames = os.listdir("picture_dir")
img_iter = cycle([cv2.imread(os.sep.join(["picture_dir",x])) for x in filenames])
key = 0
while key!=27:
    cv2.imshow('picture_windows',next(img_iter))
    key = cv2.waitKey(42)      

二、标注小工具(代碼有問題)

"""
    标注資訊格式如下:
    ('file_name',(121,232),(656,675)) # 依次:檔案名,左上角坐标,右下角坐标
"""

import os
import cv2
from tkinter.filedialog import askdirectory        # python GUI
from tkinter.messagebox import askyesno

WINDOW_NAME = 'simple Bounding Box Labeling Tool'  # 預設視窗名
FPS = 24                                           # 畫面幀率
SUPPOTED_FORMATS = ['jpg','jpeg','png']            # 支援圖像格式
DEFAULT_COLOR = {"Object":(255,0,0)}               # 物體框顔色
COLOR_GRAY = (192,192,192)                         # 資訊顯示背景和未定義物體框顯示
BAR_HEIGHT  = 16

# 上下左右,ESC,Delete鍵的cv2.waitKey() 函數傳回值
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
        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):
        if event == cv2.EVENT_LBUTTONDOWN:   # 按下左鍵,左上角坐标開始畫框,drawing狀态為True
            self._drawing = True
            self._pt0 = (x,y)
        elif event == cv2.EVENT_LBUTTONUP:   # 松開左鍵,标記右下角坐标并儲存,drawing狀态為False
            self._drawing = False
            self._pt1 = (x,y)
            self._bboxes.append((self._cur_label,self._pt0,self._pt1))  #添加标注資訊tuple
        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 = []

    # 定義标注框和目前資訊函數:在圖像下方多出 BAR_HEIGHT區域用于顯示檔案名和标注資訊
    def _draw_bbox(self,img):
        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,bpt1[1]+15),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),0.5,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 file:
                for bbox in bboxes:
                    line = repr(bbox)+"\n"
                    file.write(line)
        elif os.path.exists(filepath):
            os.remove(filepath)

    # 利用eval() 函數讀取标注框字元串資料
    @staticmethod
    def load_bbox(filepath):
        bboxes = []
        with open(filepath ,"r") as file:
            line = file.readline().rstrip()
            while line:
                bboxes.append(eval(line))
                line = file.readline().rstrip()
        return bboxes

    # 利用eval() 函數讀取物體及顔色資訊到資料
    @staticmethod
    def load_labels(filepath):
        label_colors = []
        with open(filepath ,"r") as file:
            line = file.readline().rstrip()
            while line:
                label,color = eval(line)
                label_colors[label] = color
                line = file.readline().rstrip()
        return label_colors

    # 讀取圖像檔案和對應标注框資訊(如果有的話)
    @staticmethod
    def load_sample(filepath):
        img = cv2.imread(filepath)
        bbox_filepath = get_bbox_name(filepath)
        bboxes = []
        if os.path.exists(bbox_filepath):
            bboxes = SimpleBBoxLabeling.load_bbox(bbox_filepath)
        return img,bboxes

    # 導出目前标注框資訊并清空
    def _export_n_clean_bbox(self):
        bbox_filepath = os.sep.join([self._data_dir,get_bbox_name(self._filelist[self._index])])
        self.export_bbox(bbox_filepath,self._bboxes)
        self._clean_bbox()
    # 删除目前樣本和對應的标注資訊
    def __delete_current_sample(self):
        filename = self._filelist[self._index]
        filepath = os.sep.join([self._data_dir,filename])
        if os.path.exists(filepath):
            os.remove(filepath)
        filepath = get_bbox_name(filepath)
        if os.path.exists(filepath):
            os.remove(filepath)
        self._filelist.pop(self._index)
        print("{} is deleted!".format(filename))

    # OpenCV 循環視窗:主程式
    def start(self):
        last_filename = ''                   # 最近标注檔案名,判斷是否執行圖像讀取
        label_index = 0                      # 标注物體索引
        labels = self.label_colors.keys()    # 所有标注物體名稱的清單
        n_labels = len(labels)               # 待标注物體種類數

        # 定義視窗和滑鼠回調
        cv2.namedWindow(self.window_name)
        cv2.setMouseCallback(self.window_name,self.__mouse_ops)
        key = KEY_EMPTY                      # 預設0
        delay = int(1000/FPS)                # 每次循環的持續時間
        while key!=KEY_ESC:                  # 沒按下ESC,則持續循環
            # 上下方向鍵用于選擇目前标注物體
            if key == KEY_UP:
                if label_index == 0:
                    pass
                else:
                    label_index -=1
            elif key ==KEY_DOWN:
                if label_index == n_labels-1:
                    pass
                else:
                    label_index += 1
            # 左右方向鍵用于選擇目前标注物體:
            elif key == KEY_LEFT:
                if self._index > 0:          # 首張圖檔無需清空上一張
                    self._export_n_clean_bbox()
                self._index -=1
                if self._index < 0:
                    self._index = 0
            elif key ==KEY_RIGHT:            # 末張圖檔無需清空上一張
                if self._index < len(self._filelist)-1:
                    self._export_n_clean_bbox()
                self._index +=1
                if self._index > len(self._filelist) - 1:
                    self._index = len(self._filelist) -1
            # 删除目前圖檔和對應的标注資訊
            elif key == KEY_DELETE:
                if askyesno("Delete Sample","Are you Sure?"):
                    self.__delete_current_sample()
                    key = KEY_EMPTY
                    continue

            # 如果鍵盤操作執行了換圖檔,則重新讀取,更新圖檔
            filename = self._filelist[self._index]
            if filename != last_filename:
                filepath = os.sep.join([self._data_dir,filename])
                img,self._bboxes = self.load_sample(filepath)

            # 更新目前标注物體名稱
            self._cur_label = labels[label_index]

            # 把标注和相關資訊畫在圖檔上并顯示指定時間
            canvas = self._draw_bbox(img)
            cv2.imshow(self.window_name,canvas)
            key = cv2.waitKey(delay)

            # 目前檔案名就是下次循環的老檔案名
            last_filename = filename
        print("finished!")
        cv2.destroyAllWindows()

        # 如果推出程式,需要對目前檔案進行儲存
        self.export_bbox(os.sep.join([self._data_dir,get_bbox_name(filename)]),self._bboxes)
        print("labels updatad!")
if __name__ == "__main__":
    dir_with_images = askdirectory(title="where are the images?")
    labeling_task = SimpleBBoxLabeling(dir_with_images)
    labeling_task.start()