提取 DOTA数据集中的部分类别

import os
import shutil


dota_img_path = "./DOTA/DOTA-v1.5/train/images/images/images/"
dota_label_path = "./DOTA/DOTA-v1.5/train/labelTxt-v1.5/"
save_img_path = "./FAIR1M/train/images-f_d/"
save_label_path = "./FAIR1M/train/labelTxt-f_d/"
extra_cls = ["baseball-diamond", "tennis-court", "basketball-court", "small-vehicle", "roundabout"]
# file_path = os.getcwd()
dota_label_filelist = os.listdir(dota_label_path)  # 获得DOTA数据集的所有标签文件
# dota_imgs_filelist = os.listdir(dota_img_path)  # 获得DOTA数据集的所有图片文件
count = [0, 0, 0, 0, 0]


def rename_cls(info):
    cls_name = info[-2]
    if cls_name == "baseball-diamond":
        info[-2] = "Baseball-Field"
        count[0] += 1
    if cls_name == "tennis-court":
        info[-2] = "Tennis-Court"
        count[1] += 1
    if cls_name == "basketball-court":
        info[-2] = "Basketball-Court"
        count[2] += 1
    if cls_name == "small-vehicle":
        info[-2] = "Small-Car"
        count[3] += 1
    if cls_name == "roundabout":
        info[-2] = "Roundabout"
        count[4] += 1
    return " ".join(info)  # 转换成字符串进行返回


for dota_label_file in dota_label_filelist:
    print("-------------------------------------------doing", dota_label_file, "file-------------------------------------------")
    if not dota_label_file.endswith(".txt"):
        continue
    with open(os.path.join(dota_label_path, dota_label_file)) as f:
        info = f.readline()
        crate_file = False  # 默认当前文件中不包含要提取的label
        while info:
            info = info.split(" ")
            # print(info)
            if len(info) != 10:  # 过滤掉不是标签的数据
                info = f.readline()
                continue
            cls_name = info[-2]  # 每行的倒数第二个为类别名
            if cls_name not in extra_cls:  # 如果不是要提取的类别，直接对下一行进行判断
                info = f.readline()
                continue
            # 如果程序已经走到这里，表明至少有一个能保存的label，这时应创建一个存放label的TXT文件
            if not crate_file: 
                crate_file = True   # 设置标记，也是为了防止重复创建文件
                f_out = open(os.path.join(save_label_path, dota_label_file), "a+")
            info = rename_cls(info)  # 这里返回的是字符串形式，详见rename_cls方法
            f_out.write(info)
            info = f.readline()
        if crate_file:  # 如果创建过标签文件，就需要将向对应的图片文件移动至save_img_path
            shutil.copyfile(os.path.join(dota_img_path, dota_label_file[0:-4] + ".png"), os.path.join(save_img_path, dota_label_file[0:-4] + ".png"))
        if f_out:
            f_out.close()
    # count += 1
    # if count == 50:
    #     break
for i in range(5):
    print(extra_cls[i], "=======", count[i])