深度学习目标检测图片标记工具

亦凉 2022-03-02 13:14 234阅读 0赞

import os
    import cv2
    import sys
    import argparse
    import numpy as np
    
    """ 1.遍历根目录下的图片,如果已经绘制完毕则加载其ground truth画出 2.显示一张图片，while让动作持续 3.左键点击切换至画图形模式（drawing = True），点击四下形成一个四边形，画出图形并保存四个点坐标到临时列表 4.画完任意一个多边形的时候（drawing = False），点击左键如果点击位置再某个四边形内，可将该四边形删除，更新临时列表 5.点击空格键完成当前图片编辑，保存临时为该图ground truth，切换至下一张 """
    
    KEY_EMPTY=0
    KEY_SPACE = 32  # 空格键完成本图编辑切换至下一张
    KEY_ESC=27
    extensions = ['jpg','jepg','png']
    
    
    parser = argparse.ArgumentParser(description="your script description")
    parser.add_argument('--browse', '-b', action='store_true', help='browse mode') 
    
    
    class mylabel:
        
        def __init__(self,filepath,labelpath):
            self.filepath = filepath  # 图片路径
            self.labelpath = labelpath
            self.vertex_num = 0  # 这是第0个坐标
            self.gt_boxes = []  # 所有box
            self.temp_box = [[[-1,-1],[-1,-1],[-1,-1],[-1,-1]]]  # 一个box
            self.temp_pt = (-1,-1)
        
        
        def _isRayIntersectsSegment(self,poi,s_poi,e_poi):  #[x,y] [lng,lat]
            """ 判断射线与边界的交点情况 :param poi: 需要判断的点，[x,y] :param s_poi: 边界起点 :param e_poi: 边界终点 :return: 有交点则返回True，无则返回False """
            # 输入：判断点，边起点，边终点，都是[lng,lat]格式数组
            if s_poi[1]==e_poi[1]: #排除与射线平行、重合，线段首尾端点重合的情况
                return False
            if s_poi[1]>poi[1] and e_poi[1]>poi[1]:  # 线段在射线上边
                return False
            if s_poi[1]<poi[1] and e_poi[1]<poi[1]:  # 线段在射线下边
                return False
            if s_poi[1]==poi[1] and e_poi[1]>poi[1]:  # 交点为下端点，对应spoint
                return False
            if e_poi[1]==poi[1] and s_poi[1]>poi[1]:  # 交点为下端点，对应epoint
                return False
            if s_poi[0]<poi[0] and e_poi[1]<poi[1]:  # 线段在射线左边
                return False
    
            xseg=e_poi[0]-(e_poi[0]-s_poi[0])*(e_poi[1]-poi[1])/(e_poi[1]-s_poi[1])  # 求交
            if xseg<poi[0]:  # 交点在射线起点的左侧
                return False
            return True  # 排除上述情况之后
    
        
        def _isPoiWithinPoly(self,poi,box):
            """ 判断射线与box所有边界的交点情况 :param poi: 需要判断的点，[x,y] :param box: 一个多边形 :return: 点在多边形内则返回True，否则返回False """
            sinsc=0 #交点个数
            for i in range(4):
                start_index = i
                end_index = (i + 1)%4
                start_poi = box[start_index]
                end_poi = box[end_index]
                if self._isRayIntersectsSegment(poi,start_poi,end_poi):
                    sinsc+=1 # 有交点就加1
            return True if sinsc%2==1 else  False
        
        
        def _drawLine(self,img):
            vertex_num = self.vertex_num % 4
            box = self.temp_box[0]
            if vertex_num<1:
                return img
            for i in range(vertex_num-1):
                start_vertex = (box[i][0],box[i][1])
                end_vertex = (box[i+1][0],box[i+1][1])
                cv2.line(img,start_vertex,end_vertex,(0, 255, 0),2)
            move_start_vertex = (box[vertex_num-1][0],box[vertex_num-1][1])
            move_end_vertex = self.temp_pt
            cv2.line(img,move_start_vertex,move_end_vertex,(0, 255, 0),2)
            return img
                
                
        
        def _drawBox(self,img):
            """ 绘制所有box :param img: 图像副本作为画布 :return img: 绘制完box的图像副本 """
            gt_boxes = np.array(self.gt_boxes)
            for box in gt_boxes:
                cv2.polylines(img, box, True, (0,0,255),2)
            return img
        
        
        def _mouseHandler(self, event, x, y, flags, param):
            """ 鼠标响应事件 """
            if event == cv2.EVENT_LBUTTONDOWN:  # 点击左键的时候绘四边形
                vertex_num = self.vertex_num % 4
                self.temp_box[0][vertex_num][0] = x
                self.temp_box[0][vertex_num][1]=y
                if vertex_num == 3:  # 画完一个四边形了
                    self.gt_boxes.append(self.temp_box)
                    self.temp_box = [[[-1,-1],[-1,-1],[-1,-1],[-1,-1]]]
                self.vertex_num += 1
            elif event == cv2.EVENT_MOUSEMOVE:
                self.temp_pt = (x,y)
            elif event == cv2.EVENT_RBUTTONUP:  # 点击右键的时候删除四边形，判断一个点在四边形中
                for box in self.gt_boxes:  # 遍历所有box，判断点是否在其内，在的话删除该box
                    if self._isPoiWithinPoly([x,y],box[0]):
                        self.gt_boxes.remove(box)
                        
                        
        def _saveBox(self):
            with open(self.labelpath,'w') as txt:
                contents = ""
                for box in self.gt_boxes:
                    box = box[0]
                    i = 0
                    for vertex in box:
                        contents += str(vertex[0]) + ","
                        if i!=3:
                            contents += str(vertex[1]) + ","
                            i += 1
                        else:
                            contents += str(vertex[1])
                    contents+="\n"
                txt.write(contents)
        
        
        def browseBox(self):
            img = cv2.imread(self.filepath)
            cv2.namedWindow(self.filepath)
            with open(self.labelpath,'r') as label:
                boxes = label.readlines()
                gt_boxes = []  # 所有boxes
                for box in boxes:
                    box = box.strip("\n")
                    box = box.split(",")
                    pt1 = [int(box[0]),int(box[1])]
                    pt2 = [int(box[2]),int(box[3])]
                    pt3 = [int(box[4]),int(box[5])]
                    pt4 = [int(box[6]),int(box[7])]
                    gt_box=[pt1,pt2,pt3,pt4]
                    gt_boxes.append(gt_box)
                self.gt_boxes = [gt_boxes]
                canvas = self._drawBox(img)
                cv2.imshow(self.filepath, canvas)
                key = cv2.waitKey()
                if key == KEY_SPACE:
                    cv2.destroyAllWindows()
    
    
        def label(self):
            img = cv2.imread(self.filepath)
            cv2.namedWindow("mylabel")
            cv2.setMouseCallback("mylabel", self._mouseHandler)
            key = KEY_EMPTY
            delay = int(1000/20)
    
            while key != KEY_SPACE:
                if key == KEY_ESC:  # ESC退出整个程序
                    cv2.destroyAllWindows()
                    sys.exit()
                canvas = img.copy()  # 获取图像副本
                canvas = self._drawBox(canvas)
                canvas = self._drawLine(canvas)
                cv2.imshow("mylabel",canvas)
                key = cv2.waitKey(delay)
            
            if key == KEY_SPACE:
                self._saveBox()
                cv2.destroyAllWindows()
    
        
                
    if __name__ == "__main__":
        args = parser.parse_args()
        root_dir = ".\samples"
        files = os.listdir(root_dir)
        for file in files:
            extension = file.split(".")[-1]
            if extension not in extensions:
                continue
            filename = file.split(".")[0]
            filepath = os.path.join(root_dir,file)
            labelpath = os.path.join(root_dir,filename+".txt")
            is_labeled = True if os.path.exists(labelpath) else False  # 判断文件是否已经标记
            if args.browse and is_labeled:
                print(filepath)
                ml = mylabel(filepath,labelpath)
                ml.browseBox()
            elif is_labeled:
                continue
            else:
                ml = mylabel(filepath,labelpath)
                ml.label()