Faster RCNN从demo到训练自己的数据(3)——数据集制作篇

110 阅读 0 评论 73 点赞

我是靠谱客的博主友好小丸子，最近开发中收集的这篇文章主要介绍Faster RCNN从demo到训练自己的数据(3)——数据集制作篇，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

制作数据集可以选择自己编写一个demo，也可以直接下载labelImg进行标注。Faster RCNN需要的标注文件为xml文件。

我是自己编写的一个demo，标记出每一张图片的坐标框和分类，生成的文本为txt文件，也是常见的caffe标注规格，如下图。

1.将txt文件转换为xml文件。

建议将xml文件保存在.py-faster-rcnn-masterdataVOCdevkit2007VOC2007路径下，若没有则新建。

# -*- coding:UTF-8 -*-
import os, sys
import glob
import cv2
from PIL import Image

src_img_dir = "D:/py-faster-rcnn-master/picture"

src_xml_dir = "D:/py-faster-rcnn-master/data/VOCdevkit2007/VOC2007/Annotations"

obj =[]
if os.path.exists(src_img_dir):
    list = os.listdir(src_img_dir)
    for i in range(len(list)):
        if 'txt' in list[i]:
            obj.append(list[i])


for i in range(len(obj)):
    pic = cv2.imread(src_img_dir+'/'+obj[i].strip('.txt')+'.jpg')
    size = pic.shape
    label = []
    with open(src_img_dir+'/'+obj[i],'r') as txt:
        info = true
        while info:
            info = txt.readline().strip()
            if "x00" in info:
                label.append(info)
            elif " "  in info:
                label.append(info)

    txt.close()
    if label:
        xml_file = open((src_xml_dir + '/' + obj[i].strip('.txt') + '.xml'), 'w')
        xml_file.write('<annotation>n')
        xml_file.write('    <folder>VOC2007</folder>n')
        xml_file.write('    <filename>' + obj[i].strip('.txt')+'.jpg' + '</filename>n')
        xml_file.write('    <size>n')
        xml_file.write('        <width>' + str(size[1]) + '</width>n')
        xml_file.write('        <height>' + str(size[0]) + '</height>n')
        xml_file.write('        <depth>3</depth>n')
        xml_file.write('    </size>n')

        # write the region of image on xml file
        for img_each_label in label:
            spt = img_each_label.split('x00')  # 这里如果txt里面是以逗号‘，’隔开的，那么就改为spt = img_each_label.split(',')。
            if len(spt) == 1:
                spt = img_each_label.split(' ')
            if len(spt) == 1:
                break
            if int(spt[5]) == 0 or 1 or 2:
                la = 0
            if int(spt[5]) == 3 or 4 or 5:
                la = 1
                print("yes")
            xml_file.write('    <object>n')
            xml_file.write('        <name>' + str(la) + '</name>n')
            xml_file.write('        <pose>Unspecified</pose>n')
            xml_file.write('        <truncated>0</truncated>n')
            xml_file.write('        <difficult>0</difficult>n')
            xml_file.write('        <bndbox>n')
            xml_file.write('            <xmin>' + str(spt[1]) + '</xmin>n')
            xml_file.write('            <ymin>' + str(spt[2]) + '</ymin>n')
            xml_file.write('            <xmax>' + str(spt[3]) + '</xmax>n')
            xml_file.write('            <ymax>' + str(spt[4]) + '</ymax>n')
            xml_file.write('        </bndbox>n')
            xml_file.write('    </object>n')

        xml_file.write('</annotation>')
        xml_file.close()
    else:
        print(obj[i])

2.生成train和test文件。

在.py-faster-rcnn-masterdataVOCdevkit2007VOC2007ImageSetsMain路径下生成test.txt，train.txt，trainval.txt，val.txt四个文件。test.txt是测试集，取出的占样本总量的50％；train.txt是训练集，占25％；val.txt是验证集，占25%；trainval.txt是训练和验证集，占50％。

import os
import random

trainval_percent = 0.66
train_percent = 0.5
xmlfilepath = 'D:/py-faster-rcnn-master/VOC2007_/Annotations'
txtsavepath = 'D:/py-faster-rcnn-master/VOC2007_/ImageSets/Main'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open('D:/py-faster-rcnn-master/VOC2007_/ImageSets/Main/trainval.txt', 'w')
ftest = open('D:/py-faster-rcnn-master/VOC2007_/ImageSets/Main/test.txt', 'w')
ftrain = open('D:/py-faster-rcnn-master/VOC2007_/ImageSets/Main/train.txt', 'w')
fval = open('D:/py-faster-rcnn-master/VOC2007_/ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + 'n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()