本质就是建立三个文件夹,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
需要按照自己的数据集自行更改。