PASCAL VOC提取出特定的目标类别

import os
import shutil

ann_filepath = 'D:\dataset\cityscapes\cityscape_voc_clean\Annotations\\'
img_filepath = 'D:\dataset\cityscapes\cityscape_voc_clean\JPEGImages\\'
img_savepath = 'D:\dataset\cityscapes\cityscape_voc_clean\JPEGImages_car\\'
ann_savepath = 'D:\dataset\cityscapes\cityscape_voc_clean\Annotations_car\\'
if not os.path.exists(img_savepath):
    os.mkdir(img_savepath)

if not os.path.exists(ann_savepath):
    os.mkdir(ann_savepath)
names = locals()
classes = ['pedestrian','car']

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 = '        <name>car</name>\n'

    # 在xml中找到object块，并将其记录下来
    while "    <object>\n" in lines_id_start:
        a = lines_id_start.index("    <object>\n")
        ind_start.append(a)
        lines_id_start[a] = "delete"

    while "    </object>\n" in lines_id_end:
        b = lines_id_end.index("    </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]]
    # print(string_start)
    # 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]
    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] + ".png"
        shutil.copy(name_img, img_savepath)
    fp.close()

'       <name>car</name>'