1. 程式人生 > >Faster RCNN從demo到訓練自己的資料(3)——資料集製作篇

Faster RCNN從demo到訓練自己的資料(3)——資料集製作篇

製作資料集可以選擇自己編寫一個demo,也可以直接下載labelImg進行標註。Faster RCNN需要的標註檔案為xml檔案。

我是自己編寫的一個demo,標記出每一張圖片的座標框和分類,生成的文字為txt檔案,也是常見的caffe標註規格,如下圖。

1.將txt檔案轉換為xml檔案。

建議將xml檔案儲存在.\py-faster-rcnn-master\data\VOCdevkit2007\VOC2007\路徑下,若沒有則新建。

# -*- 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-master\data\VOCdevkit2007\VOC2007\ImageSets\Main路徑下生成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()

3.將.jpg檔案放在JPEGImages資料夾下。

.\py-faster-rcnn-master\data\VOCdevkit2007\VOC2007路徑下有三個資料夾,Annotations、ImageSets和JPEGImages。