一次將自己的數據集制作成PASCAL VOC格式的慘痛經歷

因為准備訓練keras-yolo3，開源代碼上給出了voc_annotation.py文件，只要將自己的數據格式處理成PASCAL VOC格式，那么運行voc_annotation.py就可以將自己的數據集處理成模型需要的數據集。

現在我的標注數據格式如下（CSV文件，第一列是文件名，第二列對應bbox）：

圖片是文件：

不管如何先寫一個讀寫CSV文件的腳本utils.py：

# -*- coding: utf-8 -*-
# @Author  : matthew
# @File    : utils.py
# @Software: PyCharm

import csv
import os

def read_csv(csv_path, pre_dir):
    '''
    :param csv_path:csv文件路徑 
    :param pre_dir: 圖片數據所在的文件夾
    :return: 
    '''
    label_dict = {}
    with open(csv_path, "r") as f:
        reader = csv.reader(f)
        header = True
        for line in reader:
            # 除去文件頭
            if header:
                header = False
                continue
            # 處理文件存儲路徑，當做標簽
            image_path = os.path.join(pre_dir, line[0])
            # 處理后面的bbox
            bbox = []

            if line[1] is not None and len(line[1].strip()) > 0:
                for i in line[1].split(';'):
                    if i is not None and len(i.strip()) > 0:
                        bbox.append(list(map(lambda x: round(float(x.strip())), i.split('_'))))
            # 添加到label_dict
            label_dict.setdefault(image_path, bbox)
    return label_dict


def write_csv(result_dict, out_path='out.csv'):
    '''
    :param result_dict: 只一個圖片路徑，對應存儲相應bbox的list的字典
    :param out_path: 
    :return: 
    '''
    with open(out_path, 'w', newline='') as f:
        writer = csv.writer(f)
        # 寫文件頭
        writer.writerow(['name', 'coordinate'])

        for image in result_dict.keys():
            image_name = os.path.split(image)[-1]
            bbox = result_dict.get(image, [])
            bbox_rs = ';'.join(['_'.join(str(int(id)) for id in i) for i in bbox])
            writer.writerow([image_name, bbox_rs])


if __name__ == '__main__':
    label_dict = utils.read_csv(csv_path=r'./train_b.csv',
                             pre_dir=r'/home/matthew/dataset')
    write_csv(label_dict)

下面開始正式制作數據集！

第一步：改名

VOC標准數據集中圖片名稱是“000001.jpg”都為6-9位數字，jpg格式的。

因為是第一次制作，以為這種命名是必須的，我看網上也有不少人說要改名的問題。（制作完畢后，發現這一步是大可不必的，所以想省就省了吧。）

為了規范，還是進行了改名操作。

# -*- coding: utf-8 -*-
# @Author  : matthew
# @File    : pack2voc.py
# @Software: PyCharm

import os
import utils

def rename_image(label_dict={}, out_file='rename_train_b.csv'):
    '''
    改文件名的同時，修改標簽文件。並存儲成新的CSV文件rename_train_b.csv
    :param label_dict:
    :param out_file:
    :return:
    '''
    new_label_dict = {}
    i = 1
    with open(out_file, 'w') as f:
        for key in label_dict.keys():
            if not os.path.isfile(key):
                continue

            image_name = os.path.split(key)[-1]
            new_image_name = '%09d' % i + '.jpg'
            i = i + 1
            # 改名
            new_key = key.replace(image_name, new_image_name)
            os.renames(key, new_key)

            new_label_dict.setdefault(new_key, label_dict.get(key, []))
            utils.write_csv(new_label_dict, out_path=out_file)

    return out_file
 
if __name__ == '__main__':
    label_dict = utils.read_csv(csv_path=r'./train_b.csv',
                             pre_dir=r'/home/matthew/dataset')
    rename_image(label_dict)

第二步：建立VOC2007目錄

首先了解一下VOC的目錄結構：

--VOC2007
   --Annotations
   --ImageSets
     --Main
     --Layout
     --Segmentation
   --JPEGImages
   --SegmentationClass
   --SegmentationObject

Annotations  中主要存放xml文件，每一個xml對應一張圖像，
           並且每個xml中存放的是標記的各個目標的位置和類別信息，命名通常與對應的原始圖像一樣
JPEGImages  自己的原始圖像放在JPEGImages文件夾
ImageSets
           Layout	存放人體部位的數據。（用不上）
	          Main 	存放的是目標識別的數據，主要有test.txt , train.txt, val.txt，
					trainval.txt四個文件。
	          Segmentation	存放分割的數據。（用不上）

寫了個腳本生成這些文件夾：

def make_voc_dir():
    os.makedirs('VOC2007/Annotations')
    os.makedirs('VOC2007/ImageSets')
    os.makedirs('VOC2007/ImageSets/Main')
    os.makedirs('VOC2007/ImageSets/Layout')
    os.makedirs('VOC2007/ImageSets/Segmentation')
    os.makedirs('VOC2007/JPEGImages')
    os.makedirs('VOC2007/SegmentationClass')
    os.makedirs('VOC2007/SegmentationObject')
if __name__ == '__main__':
    make_voc_dir()

同時，將所有的原始圖片文件（這時候已經是命名成這種‘000000001.jpg’）移動到'JPEGImages’目錄下。

第三步：生成相應的Annotations目錄下的XML文件

首先，一份標准的VOC標注XML，格式如下：

<annotation>
	<folder>VOC2012</folder>                           
	<filename>2007_000392.jpg</filename>              //文件名
	<source>                                         //圖像來源（不重要）
		<database>The VOC2007 Database</database>
		<annotation>PASCAL VOC2007</annotation>
		<image>flickr</image>
	</source>
	<size>					      //圖像尺寸（長寬以及通道數）						
		<width>500</width>
		<height>332</height>
		<depth>3</depth>
	</size>
	<segmented>1</segmented>		//是否用於分割（在圖像物體識別中01無所謂）
	<object>                        //檢測到的物體
		<name>horse</name>          //物體類別
		<pose>Right</pose>          //拍攝角度
		<truncated>0</truncated>    //是否被截斷（0表示完整）
		<difficult>0</difficult>    //目標是否難以識別（0表示容易識別）
		<bndbox>                    //bounding-box（包含左下角和右上角xy坐標）
			<xmin>100</xmin>
			<ymin>96</ymin>
			<xmax>355</xmax>
			<ymax>324</ymax>
		</bndbox>
	</object>
	<object>                        //檢測到多個物體
		<name>person</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>198</xmin>
			<ymin>58</ymin>
			<xmax>286</xmax>
			<ymax>197</ymax>
		</bndbox>
	</object>
</annotation>

我們的主要任務就是將CSV中的每一行數據轉換成這種格式，然而網上找到的最多的竟然是一堆matlab的代碼。什么鬼？！所以，只好自己動手擼代碼。

def save_xml(image_name, bbox, save_dir='./VOC2007/Annotations', width=1609, height=500, channel=3):
      '''
    將CSV中的一行
    000000001.jpg [[1,2,3,4],...]
    轉化成
    000000001.xml

    :param image_name:圖片名
    :param bbox:對應的bbox
    :param save_dir:
    :param width:這個是圖片的寬度，博主使用的數據集是固定的大小的，所以設置默認
    :param height:這個是圖片的高度，博主使用的數據集是固定的大小的，所以設置默認
    :param channel:這個是圖片的通道，博主使用的數據集是固定的大小的，所以設置默認
    :return:
    '''
    from lxml.etree import Element, SubElement, tostring
    from xml.dom.minidom import parseString

    node_root = Element('annotation')

    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'JPEGImages'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image_name

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = '%s' % width

    node_height = SubElement(node_size, 'height')
    node_height.text = '%s' % height

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '%s' % channel

    for x, y, w, h in bbox:
        left, top, right, bottom = x, y, x + w, y + h
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = 'car'
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = '%s' % left
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = '%s' % top
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = '%s' % right
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = '%s' % bottom

    xml = tostring(node_root, pretty_print=True)  
    dom = parseString(xml)

    save_xml = os.path.join(save_dir, image_name.replace('jpg', 'xml'))
    with open(save_xml, 'wb') as f:
        f.write(xml)

    return


def change2xml(label_dict={}):
    for image in label_dict.keys():
        image_name = os.path.split(image)[-1]
        bbox = label_dict.get(image, [])
        save_xml(image_name, bbox)
    return


if __name__ == '__main__':
    # step 2
    # make_voc_dir()

    # step 3
    # label_dict = utils.read_csv(csv_path=r'./train_b.csv',
    #                             pre_dir=r'/home/matthew/dataset')
    # rename_image(label_dict)

    # step 3
    label_dict = utils.read_csv(csv_path=r'./rename_train_b.csv',
                                pre_dir=r'/home/matthew/VOC2007/JPEGImages')
    change2xml(label_dict)

第四步：生成Main目錄下的txt文件

這一步很簡單，就是生成測試、驗證數據集合等等，然后存儲成txt文件，網上有博主提供了代碼，照搬如下。

建立make_train_val_test_set.py，放在VOC2007目錄下，然后運行。

# -*- coding: utf-8 -*-
# @Author  : matthew
# @File    : make_train_val_test_set.py
# @Software: PyCharm

import os
import random


def _main():
    trainval_percent = 0.1
    train_percent = 0.9
    xmlfilepath = 'Annotations'
    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('ImageSets/Main/trainval.txt', 'w')
    ftest = open('ImageSets/Main/test.txt', 'w')
    ftrain = open('ImageSets/Main/train.txt', 'w')
    fval = open('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:
                ftest.write(name)
            else:
                fval.write(name)
        else:
            ftrain.write(name)

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


if __name__ == '__main__':
    _main()

第五步：運行voc_annotation.py

運行的時候，注意修改這個腳本里面的一些路徑和參數。

import xml.etree.ElementTree as ET
from os import getcwd

# 注意這里的‘2007’，也許你的就需要修改
sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]

# 注意類別
classes = ["car"]
def convert_annotation(year, image_id, list_file):
    # 注意路徑
    in_file = open('VOC%s/Annotations/%s.xml'%(year, image_id))
    tree=ET.parse(in_file)
    root = tree.getroot()

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))

wd = getcwd()

for year, image_set in sets:
    # 注意路徑
    image_ids = open('VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        # 注意路徑
        list_file.write('%s/VOC%s/JPEGImages/%s.jpg'%(wd, year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()
    

后記

成功之后，會生成三個文件，像這樣

打開看一下：

/home/matthew/VOC2007/JPEGImages/000000160.jpg 186,192,353,349,0 579,286,850,500,0
/home/matthew/VOC2007/JPEGImages/000000162.jpg 403,22,458,60,0 400,245,552,389,0 432,0,459,12,0 926,1,999,15,0
/home/matthew/VOC2007/JPEGImages/000000166.jpg 146,246,340,428,0

呵呵，有沒有很失望！！！

原來需要的這種格式的文件：

文件的絕對路徑 left,top,right,bottom,類別編號

一開始就知道的話，分分鍾的事就可以完成轉化，結果折騰了一圈~2333333333333

寫博不易，喜歡請打賞。

主要參考

https://blog.csdn.net/Patrick_Lxc/article/details/80615433

https://blog.csdn.net/lilai619/article/details/79695109

https://blog.csdn.net/zhangjunbob/article/details/52769381