VOC2012数据集中选出自己想要的数据集

""" Created on Thu Sep 3 08:33:42 2020 @author: Administrator """ import os import shutil ann_filepath='D:/dataset/VOC/train/VOCdevkit/VOC2012/Annotations/' img_filepath='D:/dataset/VOC/train/VOCdevkit/VOC2012/JPEGImages/' img_savepath='D:/dataset/VOC/train/VOCdevkit/VOC2012/JPEGImages_ssd/' ann_savepath='D:/dataset/VOC/train/VOCdevkit/VOC2012/Annotations_ssd/' def voc_selcet(ann_filepath,img_filepath,img_savepath,ann_savepath): names = locals() classes = ['aeroplane','bicycle','bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow','diningtable', 'dog', 'horse', 'motorbike', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor', 'person'] if not os.path.exists(img_savepath): os.mkdir(img_savepath) if not os.path.exists(ann_savepath): os.mkdir(ann_savepath) for file in os.listdir(ann_filepath): print(file) fp = open(ann_filepath + '/' + file) ann_savefile=ann_savepath+file fp_w = open(ann_savefile, 'w') lines = fp.readlines() ind_start = [] ind_end = [] lines_id_start = lines[:] lines_id_end = lines[:] classes1 = '\t\t<name>bicycle</name>\n' classes2 = '\t\t<name>motorbike</name>\n' classes3 = '\t\t<name>bus</name>\n' classes4 = '\t\t<name>car</name>\n' classes5 = '\t\t<name>person</name>\n' #在xml中找到object块，并将其记录下来 while "\t<object>\n" in lines_id_start: a = lines_id_start.index("\t<object>\n") ind_start.append(a) lines_id_start[a] = "delete" while "\t</object>\n" in lines_id_end: b = lines_id_end.index("\t</object>\n") ind_end.append(b) lines_id_end[b] = "delete" #names中存放所有的object块 i = 0 for k in range(0, len(ind_start)): names['block%d' % k] = [] for j in range(0, len(classes)): if classes[j] in lines[ind_start[i] + 1]: a = ind_start[i] for o in range(ind_end[i] - ind_start[i] + 1): names['block%d' % k].append(lines[a + o]) break i += 1 #print(names['block%d' % k]) #xml头 string_start = lines[0:ind_start[0]] #xml尾 string_end = [lines[len(lines) - 1]] #在给定的类中搜索，若存在则，写入object块信息 a = 0 for k in range(0, len(ind_start)): if classes1 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes2 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes3 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes4 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] if classes5 in names['block%d' % k]: a += 1 string_start += names['block%d' % k] string_start += string_end for c in range(0, len(string_start)): fp_w.write(string_start[c]) fp_w.close() #如果没有我们寻找的模块，则删除此xml，有的话拷贝图片 if a == 0: os.remove(ann_savepath+file) else: name_img = img_filepath + os.path.splitext(file)[0] + ".jpg" shutil.copy(name_img, img_savepath) fp.close()