一次将自己的数据集制作成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