kitti数据集打包为voc数据格式

其实，这篇文章跟我的yolov4那篇文章有些重复，但是想着好歹把代码写写改了一些时间，还是写下来，纪念一下吧。哎，简单说下吧，我们将我们要训练的数据全部放在JPEGImages这个文件夹下面，然后把这些照片的对应的txt标记信息也放在一个文件夹下面，然后就可以使用下面的代码将数据集转化为voc格式的数据集了，直接看代码吧。

'''author:nike hu'''
# -*- coding: utf-8 -*-

import shutil
import os
import cv2

headstr = """\
<annotation>
    <folder>VOC2007</folder>
    <filename>%06d.jpg</filename>
    <source>
        <database>My Database</database>
        <annotation>PASCAL VOC2007</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""

tailstr = '''\
</annotation>
'''




def writexml(idx, head, bbxes, tail):
    filename = ("Annotations/%06d.xml" % (idx))
    f = open(filename, "w")
    f.write(head)
    for bbx in bbxes:
        f.write(objstr % (bbx[-1], bbx[0], bbx[1], bbx[0] + bbx[2], bbx[1] + bbx[3]))
    f.write(tail)
    f.close()


def clear_dir():
    if shutil.os.path.exists(('Annotations')):
        shutil.rmtree(('Annotations'))
    if shutil.os.path.exists(('ImageSets')):
        shutil.rmtree(('ImageSets'))
    # if shutil.os.path.exists(('JPEGImages')): # 因为我们已经将所有图片放到这个文件夹里面了，所以不需要再创建了
    #     shutil.rmtree(('JPEGImages'))

    shutil.os.mkdir(('Annotations'))
    shutil.os.makedirs(('ImageSets/Main'))
    # shutil.os.mkdir(('JPEGImages'))


def excute_datasets():
    '''在main文件夹下面要创建四个文件，trainval是总样本的百分之九十，train是总数据的百分之七十，val是总数据样本的百分之20，剩下的百分之10是测试样本'''
    ftrainval = open(('ImageSets/Main/' + 'trainval' + '.txt'), 'a')
    ftrain = open(('ImageSets/Main/' + 'train' + '.txt'), 'a')
    fval = open(('ImageSets/Main/' + 'val' + '.txt'), 'a')
    ftest = open(('ImageSets/Main/' + 'test' + '.txt'), 'a')
    images = './JPEGImages/' # 这是存储图片的位置
    txtfile = './label/' # 这个是是存储标记信息的文件夹
    txtlist = os.listdir(txtfile)
    lenfile = len(txtlist) # 这个是标记的信息的总的文件数量
    count = 1 # 统计正在处理的数量
    for txtname in txtlist:
        txt_path = os.path.join(txtfile, txtname)
        image_path = os.path.join(images, txtname.split('.')[0] + '.png')  # 这里是图片的路径
        im = cv2.imread(image_path)  # 读取图片
        if im is None:  # 如果不存在这张照片，跳过
            continue
        head = headstr % (int(txtname.split('.')[0]), im.shape[1], im.shape[0], im.shape[2]) # xml文件的头部分
        boxes = []
        with open(txt_path, 'r') as f:
            while True:
                txt_content = f.readline().split(' ')
                if txt_content[0] == '':
                    break
                label_name = txt_content[0]
                if label_name == 'Misc' or label_name == 'DontCare': # 如果是这两类就去掉
                    continue
                box = [float(x) for x in txt_content[4:8]] # 这个是坐标
                box.append(label_name) # 把每个坐标对应的标签加入
                boxes.append(box)
        writexml(int(txtname.split('.')[0]), head, boxes, tailstr)
        if count < 0.9 * lenfile: # 总样本的百分之90部分存到trainval
            ftrainval.write('%06d\n' % (int(txtname.split('.')[0])))
            if count < 0.7 * lenfile: # 总样本的百分之70存入train
                ftrain.write('%06d\n' % (int(txtname.split('.')[0])))
            else: # 在0.7到0.9之间的数据存入val文件
                fval.write('%06d\n' % (int(txtname.split('.')[0])))
        else:
            ftest.write('%06d\n' % (int(txtname.split('.')[0])))
        count += 1
    ftrain.close() # 运行的时候出现过没有存进去的情况，原因是数据在内存中，还没有存在磁盘中，一般程序运行完会将数据放到磁盘中，或者用close语句
    ftest.close()
    ftrainval.close()
    fval.close()



if __name__ == '__main__':
    clear_dir()
    idx = excute_datasets()
    print('Complete...')

就这样吧，没啥好说的，在yolov4里面已经详细说了这样做的思路。

2020 5.1