faster-rcnn之生成训练数据

目标检测的数据无论是从imagenet上获取还是通过labelImg人工标注获得，最后的标注文件都遵循VOC标注文件的格式，本文提供一种VOC标注数据到faster_rcnn训练数据的转换代码，同时代码对标注边框处理，防止进入训练由于边框坐标问题引起报错。首先看一下faster-rcnn要求的训练数据的结构图如下：

代码用到另一篇文章的工具类，请自行前往下载加载： Imagenet标注文件的Read和Write

代码如下：label_imagenet2rcnn.py

#coding=utf-8
import os,cv2
from image_label_util import PicAnno,PicObject,VocUtil
import logging
logging.basicConfig(
	level=logging.DEBUG,
	format='%(asctime)s %(levelname)s: %(message)s',
	datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)


class LabelImagenet2rcnn:
    def __init__(self, orgi_root_dir, tar_root_dir):
        self.root_dir = orgi_root_dir
        self.img_dir = os.path.join(self.root_dir, 'JPEGImages')
        self.anno_dir = os.path.join(self.root_dir, 'Annotations')

        self.voc2007_root_dir = tar_root_dir
        self.voc2007_jpegimages_dir = os.path.join(self.voc2007_root_dir, 'JPEGImages')
        self.voc2007_annotations_dir = os.path.join(self.voc2007_root_dir, 'Annotations')
        self.voc2007_imagesets_main_dir = os.path.join(self.voc2007_root_dir, 'ImageSets', 'Main')
        if not os.path.exists(self.voc2007_root_dir):
            os.mkdir(self.voc2007_root_dir)
        if not os.path.exists(self.voc2007_jpegimages_dir):
            os.mkdir(self.voc2007_jpegimages_dir)
        if not os.path.exists(self.voc2007_annotations_dir):
            os.mkdir(self.voc2007_annotations_dir)
        if not os.path.exists(self.voc2007_imagesets_main_dir):
            os.makedirs(self.voc2007_imagesets_main_dir)
        self.vocUtil = VocUtil()

    def etl_xml(self, pic_im, picAnno):
        h, w = pic_im.shape[:2]
        for picObj in picAnno.objects:
            if int(picObj.xmin) >= int(picObj.xmax) or int(picObj.ymin) >= int(picObj.ymax):
                xmin,xmax,ymin,ymax = int(picObj.xmin) ,int(picObj.xmax), int(picObj.ymin), int(picObj.ymax)
                picObj.xmin = str(xmin if xmin < xmax else xmax)
                picObj.ymin = str(ymin if ymin < ymax else ymax)
                picObj.xmax = str(xmax if xmax > xmin else xmin)
                picObj.ymax = str(ymax if ymax < ymax else ymin)

            if int(picObj.xmin) <= 0:
                picObj.xmin = '1'
            if int(picObj.ymin) <= 0:
                picObj.ymin = '1'
            if int(picObj.xmax) >= int(w):
                picObj.xmax = str(int(w) - 1)
            if int(picObj.ymax) >= int(h):
                picObj.ymax = str(int(h) - 1)
        return picAnno

    def imagenet2rcnn(self):
        file_names = [xml_name for xml_name in os.listdir(self.anno_dir) if xml_name.endswith('.xml')]
        for idx, xml_name in enumerate(file_names):
            if idx > 0 and idx % 100 == 0:
                logger.info('total num:{} done num:{}'.format(len(file_names), idx))
            xml_path = os.path.join(self.anno_dir, xml_name)
            pic_path = os.path.join(self.img_dir, xml_name[:-4] + '.jpg')
            if not os.path.exists(pic_path):
                pic_path = os.path.join(self.img_dir, xml_name[:-4] + '.png')
                if not os.path.exists(pic_path):
                    continue

            pic_im = cv2.imread(pic_path)
            picAnno = self.vocUtil.read_anno_xml(xml_path)
            picAnno = self.etl_xml(pic_im,picAnno)

            xml_text = self.vocUtil.parse_anno_xml(picAnno)
            new_xml_name = str(idx).zfill(6) + '.xml'
            new_pic_name = str(idx).zfill(6) + '.jpg'
            self.vocUtil.save_anno_xml(os.path.join(self.voc2007_annotations_dir, new_xml_name), xml_text)
            cv2.imwrite(os.path.join(self.voc2007_jpegimages_dir, new_pic_name),pic_im)
        self.vocUtil.gene_train_test_val_txt(self.voc2007_annotations_dir,self.voc2007_imagesets_main_dir)


if __name__ == '__main__':
    orgi_root_dir = r'//Users/songhongwei/data/cashier/cashier_test'
    tar_root_dir = r'/Users/songhongwei/data/cashier/VOC2007'
    imagenet2rcnn = LabelImagenet2rcnn(orgi_root_dir,tar_root_dir)
    imagenet2rcnn.imagenet2rcnn()

以五张标注图片为例：
1、生成的txt文件中每行为图片名称（不包含扩展名）

2、文件路径结构图