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