从DIOR数据集中提取自己需要的类别数据

        解压后的DIOR数据集包括四个文件夹，Annotations是xml标签文件，ImageSet是划分好的三个数据集，另两个文件夹是对应的影像。

 

        DIOR数据集涵盖20个对象类（数据论文地址：Object Detection in Optical Remote Sensing Images: A Survey and A New Benchmark），并且所有类别的图片和标注数据都是混在一个文件夹下面的，由于我只需要其中的部分类，所以想把他们单独拿出来。 

        首先是利用xml文件的<name>节点把需要类别的filename单独放在一个文件夹下面，在这篇博客的基础上浅改一下：

# from os import listdir, getcwd
# from os.path import join
import cv2
import os
from xml.etree.ElementTree import parse
# 不需要提前新建txt

xmlfilepath = 'F:\\DIOR数据集\\Annotations\\'  # 存放xml文件的文件夹路径
total_xml = os.listdir(xmlfilepath)  # 遍历xml文件
for xml in total_xml:

    xml_path = os.path.join(xmlfilepath, xml)  # xml的全路径，也就是文件的路径

    tree = parse(xml_path)  # 获取ElementTree
    root = tree.getroot()  # 获取根元素

    filename = root.findtext('filename')[0:5]  # 提取到了xml对应的filename的名字，[0:5]是我只想要名字，不想要后缀.jpg

    for obj in root.iter('object'):  # 如有很多object ，需要遍历object
        cls = obj.find('name').text  # 将解析得到该keypoints下name中的标签
        if cls == 'bridge':  # 选择自己需要的标签类（如桥bridge），然后就可以把这类别下需要的坐标等信息输出了

            '''xmlbox = obj.find('bndbox')  # 坐标信息
            b = ((xmlbox.find('xmin').text), (xmlbox.find('ymin').text),
                 (xmlbox.find('xmax').text), (xmlbox.find('ymax').text))'''
            with open('F:\\DIOR数据集\\2.txt', 'a') as f:  # txt存储路径
                f.write(str(filename))
                # print("b:", b[0])
                '''for i in range(len(b)):
                    f.write(' ')
                    f.write(str(b[i]))'''
                f.write('\n')
        break  # 因为我只要名称，所以不需要遍历object，只要有bridge就跳出最深层的for循环

         然后再利用这个txt把对应名称的jpg和xml文件放在新的文件夹下面。参考这个

# 导入需要使用的python包
import shutil

# 根据txt文件中存储的文件名，提取对应的文件保存到另一个文件夹
data = []

# 读取存储txt文件
for line in open("F:/DIOR数据集/2.txt", "r"):  # 设置文件对象并读取每一行文件
    data.append(line)

for a in data:
    # src是总文件夹
    src1 = 'F:/DIOR数据集/JPEGImages-test/{}.jpg'.format(a[:-1])
    src2 = 'F:/DIOR数据集/Annotations/{}.xml'.format(a[:-1])

    # dst是保存提取结果的文件夹
    # jpg
    #dst1 = 'F:/DIOR数据集/jpg/{}.jpg'.format(a[:-1])
    #shutil.move(src1, dst1)

    # xml
    dst2 = 'F:/DIOR数据集/xml/{}.xml'.format(a[:-1])
    shutil.move(src2, dst2)

OK啦