voc格式数据集去掉多余的类别

最新推荐文章于 2024-05-06 00:08:33 发布

weixin_43964394

最新推荐文章于 2024-05-06 00:08:33 发布

阅读量191

点赞数 1

文章标签： python 深度学习开发语言

本文链接：https://blog.csdn.net/weixin_43964394/article/details/129239473

版权

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import xml.etree.ElementTree as ET
import shutil

# 根据自己的情况修改相应的路径
ann_filepath = 'G:/DataSets/VisDrone2019/VisDrone2019-DET-val/Annotations_XML/'
img_filepath = 'G:/DataSets/VisDrone2019/VisDrone2019-DET-val/images/'
img_savepath = 'G:/DataSets/VisDrone2019/VisDrone2019-DET-val/test/JPEGImages/'
ann_savepath = 'G:/DataSets/VisDrone2019/VisDrone2019-DET-val/test/Annotations/'
if not os.path.exists(img_savepath):
    os.mkdir(img_savepath)

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

# 这是VOC数据集中所有类别
# classes = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
#             'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
#              'dog', 'horse', 'motorbike', 'pottedplant',
#           'sheep', 'sofa', 'train', 'person','tvmonitor']

classes = ['pedestrian','people','bicycle','car','van','truck','tricycle','awning-tricycle','bus','motor']  # 这里是需要提取的类别


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)
    # 文本文件名自己定义，主要用于生成相应的训练或测试的txt文件
    with open('G:/DataSets/VisDrone2019/VisDrone2019-DET-val/test/test.txt', 'a') as file_txt:
        file_txt.write(os.path.splitext(file)[0])
        file_txt.write("\n")
    return True


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

用的别人代码，简单记录一下以防遗忘