从VOC数据集选择部分类别数据

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@Author  :   {FirstElfin}
@License :   (C) Copyright 2013-2020, {DHWL}
@Contact :   {2968793701@qq.com}
@Software:   PyCharm
@File    :   test.py
@Time    :   11/22/19 11:55 AM
"""
import os
import xml.etree.ElementTree as ET
import shutil

ann_filepath = './VOC2007_22/Annotations/'
img_filepath = './VOC2007_22/JPEGImages/'
img_savepath = './VOC2007/JPEGImages/'
ann_savepath = './VOC2007/Annotations/'
if not os.path.exists(img_savepath):
    os.mkdir(img_savepath)

if not os.path.exists(ann_savepath):
    os.mkdir(ann_savepath)

classes = ['bicycle', 'bus', 'car', 'motorbike', 'train']


def save_annotation(file):

    tree = ET.parse(ann_filepath + '/' + file)
    root = tree.getroot()
    result = root.findall("object")
    bool_num = 0
    for obj in result:
        if obj.find("name").text not in classes:
            root.remove(obj)
        else:
            bool_num = 1
    if bool_num:
        tree.write(ann_savepath + file)
        return True
    else:
        return False


def save_images(file):
    name_img = img_filepath + os.path.splitext(file)[0] + ".jpg"
    shutil.copy(name_img, img_savepath)
    return True


if __name__ == '__main__':
    for f in os.listdir(ann_filepath):
        if save_annotation(f):
            save_images(f)

　　./VOC2007_22是备份的数据集(原数据集)，./VOC2007是我们自己要生成的数据集(制作自己的数据)

　　通过classes可以选择你需要的类别，原始类别为：

classes = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
           'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
           'dog', 'horse', 'motorbike', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor', 'person']