区分数据集每张图片的类别并分别保存

from cgi import print_arguments
import shutil
import random
import os
import pathlib
import numpy as np
import pandas as pd
from numpy import size
from collections import Counter
# from torch import classes

img_path = 'val/val/images'
txt_path = 'val/val/labels'
all_path = 'classes/'

size = [800, 800]
classes = ['ship', 'airplane', 'vehicle', 'wind-mill', 'storage-tank', 'bridge', 'swimming-pool', 'person']

txt_name = os.listdir(txt_path)  #列出全部的txt
img_name = os.listdir(img_path)

# all_label = np.array([])

for txt in txt_name:  
    txtpath = os.path.join(txt_path,txt)
    all_label = np.array([])
    gt = np.loadtxt(open(txtpath, encoding='utf8'),dtype=np.str, delimiter=' ')
    if(gt.ndim==1):
        gt = np.expand_dims(gt, axis=0)   
    length = gt.shape[0]
    for j in range(length):
        label_ = gt[j][4]
        all_label = np.append(all_label,str(label_))  #将一个txt的 所有类别都保存下来
    print(all_label)
    dict_ = {}
    for i in all_label:
        if i not in dict_.keys():
            dict_[i] = Counter(all_label)[i]
    print(dict_)
    class1_ = max(dict_, key = dict_.get) 
    print(class1_)
    if (class1_ in classes):
        savetxt = all_path + str(class1_) + '/txt'  #保存该txt的路径
        shutil.copy(txtpath, savetxt)
        txt = str(txt)
        img = txt[:-4] + ".jpg"
        saveimg = all_path + str(class1_) + '/img'
        if img in img_name:
            imgpath = img_path + '/' + img
            shutil.copy(imgpath, saveimg)