制作自己的VOC2007格式数据

本质就是建立三个文件夹,Annotation用于存放xml标记文件,JPEGImages用于存放image数据,ImageSets文件夹下的Main文件夹存放train.txt,val.txt,trainval.txt,test.txt四个list。代码如下:

import os

def dir2txt(dir, txt):
    files = os.listdir(dir)
    fw = open(txt, 'w')
    for file in files:
        line = file.strip().split('.')[0]+'\n'
        fw.writelines(line)
        print line
    fw.close()

def txt2xml(txt, xml):
    import cv2

    s1 = """    <object>
        <name>{0}</name>
        <pose>Unspecified</pose>
        <truncated>{1}</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{2}</xmin>
            <ymin>{3}</ymin>
            <xmax>{4}</xmax>
            <ymax>{5}</ymax>
        </bndbox>
    </object>"""
    s2 = """<annotation>
    <folder>VisDrone</folder>
    <filename>{0}</filename>
    <source>
        <database>VisDrone</database>
        <annotation>VisDrone</annotation>
        <image>VisDrone</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>VisDrone</name>
    </owner>
    <size>
        <width>{1}</width>
        <height>{2}</height>
        <depth>{3}</depth>
    </size>
    <segmented>0</segmented>{4}
</annotation>
"""

    key_value = ['ignored', 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor', 'others']

    text_ = txt
    flabel = open(text_, 'r')
    lb = flabel.readlines()
    flabel.close()
    ob2 = ""
    for i in range(len(lb)):
        y2 = lb[i].split(',')
        y3 = [int(i) for i in y2[:4]]
        y3[2] = y3[2]+y3[0]
        y3[3] = y3[3]+y3[1]
        cls = key_value[int(y2[5])]
        trn = y2[6]
        ob2 += '\n' + s1.format(cls, trn, y3[0], y3[1], y3[2], y3[3])
    imgname = text_.strip().split('.')[0].split('/')[-1] + '.jpg'
    imgpath = text_.strip().split('annotations')[0]+'images/'+imgname
    img = cv2.imread(imgpath)
    height, width, depth = img.shape
    savename = xml
    f = open(savename, 'w')
    ob1 = s2.format(imgname, width, height, depth, ob2)
    f.write(ob1)
    f.close()

if __name__ == '__main__':

    '''Write the txt'''
    # dir = '/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-val/images/'
    # txt = '/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/val.txt'
    # dir2txt(dir, txt)
    # dir = '/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-train/images/'
    # txt = '/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/train.txt'
    # dir2txt(dir, txt)
    # dir = '/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-test-challenge/images/'
    # txt = '/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/test.txt'
    # dir2txt(dir, txt)
    # fw = open('/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/trainval.txt','w')
    # f1 = open('/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/val.txt','r')
    # f2 = open('/data/yutingzhao/VisDronedevkit/VisDrone/ImageSets/Main/train.txt','r')
    # for line in f1.readlines():
    #     fw.writelines(line)
    # f1.close()
    # for line in f2.readlines():
    #     fw.writelines(line)
    # f2.close()
    # fw.close()

    '''Write the xml'''
    xml_fold = "/data/yutingzhao/VisDronedevkit/VisDrone/Annotations/"

    train_fold = os.listdir("/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-train/annotations/")
    for txt in train_fold:
        xml = xml_fold+txt.strip().split('.')[0]+'.xml'
        txt = "/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-train/annotations/"+txt
        txt2xml(txt, xml)
        print xml

    # val_fold = os.listdir("/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-val/annotations/")
    # for txt in val_fold:
    #     xml = xml_fold+txt.strip().split('.')[0]+'.xml'
    #     txt = "/data/yutingzhao/VisDronedevkit/VisDrone/VisDrone2018-DET-val/annotations/"+txt
    #     txt2xml(txt, xml)
    #     print xml
需要按照自己的数据集自行更改。

猜你喜欢

转载自blog.csdn.net/yutingzhaomeng/article/details/80270423
今日推荐