python制作单类别检测数据集

为了训练自己的yolov3，但是训练一个公开数据集需要的时间太长了，所以就随手做了一个制作数据集的小demo，代码不长比较简单，但是只能标注矩形框，而且只可以制作单类别数据集。

制作前的准备：将需要的图片放在一个文件夹下，最好都是一个格式的

生成的数据集格式：每行代表一个图像标签，格式为：图像名 x1，y1，x2，y2 …

标注过程如下：将标注的框指示在图像上，按下鼠标左键开始画框，松开鼠标左键开始画框，当标注完此图片时，按q结束此张图片开始下一张图片的标注。

代码如下：

import cv2
import numpy as np
import os

class dataset_creator(object):
    def __init__(self,image_path,txt_path):
        self.image_path = image_path    #存放图像的路径
        self.txt_path = txt_path        #标签的路径
        self.file = open(self.txt_path,'w')
        self.image_names = os.listdir(image_path)   #遍历文件夹下的文件
        self.image_name = None        #当前操作的图像名
        self.recs = []                #已经选择的bbox
        self.drawing = False          #鼠标按下，开始画框
        self.ix = -1                  #鼠标按下时的x坐标
        self.iy = -1                  #鼠标按下时的y坐标
        self.image = None             #当前操作的原图像
        self.drawn_image = None       #在原图像上画框后的图像

    def get_dataset(self):
        for image_name in self.image_names:       #遍历标签的图像名
            self.image_name = image_name
            self.image = cv2.imread(self.image_path+'/'+image_name)   #读取图像
            cv2.namedWindow(image_name)      #这个时必须要的
            cv2.setMouseCallback(image_name, self.draw_rec)   #设置鼠标事件
            while True:
                cv2.imshow(image_name,self.image)
                key = cv2.waitKey(0)   #等待键盘事件
                if key == 113:         #按下q时结束为此图片打标签并开始下一张
                    print(self.recs)
                    self.write_txt(image_name,self.recs)
                    break
            cv2.destroyWindow(image_name)    #销毁当前图像显示窗口
            self.recs = []   #清空当前图像已选择的bbox
        self.file.close()   #关闭文件对象才能保存

    def draw_rec(self,event, x, y, flags, param):
        for rec in self.recs:   #把已经确定的框画上去
            cv2.rectangle(self.image, (rec[0],rec[1]), (rec[2],rec[3]), (0, 255, 0), 2)
        if event == cv2.EVENT_LBUTTONDOWN:   #按下鼠标左键
            self.drawing = True
            self.ix, self.iy = x, y
        elif event == cv2.EVENT_MOUSEMOVE and flags == cv2.EVENT_FLAG_LBUTTON:  #移动鼠标，实时显示框
            if self.drawing is True:
                self.drawn_image = cv2.rectangle(np.copy(self.image), (self.ix, self.iy), (x, y), (0, 255, 0), 2)
                cv2.imshow(self.image_name, self.drawn_image)

        elif event == cv2.EVENT_LBUTTONUP:   #抬起鼠标左键
            self.drawing = False
            self.recs.append([self.ix,self.iy,x,y])
            #cv2.rectangle(self.image, (self.ix, self.iy), (x, y), (0, 255, 0), 2)

    def write_txt(self,image_name,recs):   #将当前图像名和属于该图像的bbox存入txt文件
        self.file.writelines(image_name+' ')
        for rec in recs:
            for xy in rec:
                self.file.writelines(str(xy)+',')
            self.file.writelines('0 ')
        self.file.writelines('\n')

if __name__ == '__main__':
    creator = dataset_creator("./images","data.txt")
    creator.get_dataset()