提取目录中的图片和文本保存到新的文件夹并划分数据集

from importlib.resources import path
import os
import random
import shutil 

#随机打乱数据集  可以不打乱txt文件
def get_list(img_path):
    return [os.path.join(img_path,f) for f in os.listdir(img_path) if f.endswith('.jpg')]

#输入划分比例train:test 7:3 p_train p_test
def get_dataset(txt_path, img_path, save_txt_path, save_img_path,p_train):
    data_img = os.listdir(img_path)
    img_length = len(data_img)
    # img_list = get_list(img_path) #打乱image
    random.shuffle(data_img)  #打乱列表元素
    print(data_img)
    exit()
    length_img = int(img_length*p_train)  #划分比例
    for n in data_img[length_img:]:
        # shutil.move(n,save_img_path)#移动到一个新的文件夹中
        shutil.copy(img_path+'/'+n,save_img_path+'/'+n)#复制文件到新的文件夹中
    select_txt(save_img_path,txt_path,save_txt_path) #选取txt文件    

#根据选择好的image选取txt
def select_txt(save_img_path,txt_path,save_txt_path):
    img_ = os.listdir(save_img_path)
    txt_ = os.listdir(txt_path)
    for i in img_:
        txt_name = i[:-4] + '.txt'
        if txt_name in txt_:
            # shutil.move(txt_name,save_txt_path)
            shutil.copy(save_txt_path+'/'+txt_name,txt_path+'/'+txt_name)

#重新rename txt and img
def rename(save_txt_path,save_img_path):
    txt_ = os.listdir(save_txt_path)
    img_ = os.listdir(save_img_path)
    i = 1
    for file in txt_:
        new_txt_name = file.replace(file,"%d.jpg"%i).zfill(10)
        print(new_txt_name)
        os.rename(os.path.join(save_txt_path,file),os.path.join(save_txt_path,new_txt_name))
        i +=1
    for file in img_:
        new_img_name = file.replace(file,"%d.jpg"%i).zfill(10)
        print(new_img_name)
        os.rename(os.path.join(save_img_path,file),os.path.join(save_img_path,new_img_name))
        i +=1
    
if __name__ == '__main__':
            txt_path = '数据集分割测试/img'  #注释文件路径txt
            img_path = '数据集分割测试/txt'  #图片路径
            save_txt_path = '数据集分割测试/save_txt' #txt保存路径
            save_img_path = '数据集分割测试/save_img' #img保存路径
            p_train = 0.7
            get_dataset(txt_path,img_path,save_txt_path,save_img_path,p_train)
            rename(save_txt_path,save_img_path)  

#数据集划分
import shutil
import random
import os
import pathlib

#car
img_path = 'car/images/'  #原始img 路径car\airplane
label_path = 'car/txt/' #原始label路径
len_img_path = os.listdir(img_path)
len_label_path = os.listdir(label_path)
print(len(len_img_path))
print(len(len_label_path))

save_train_img = 'car/save_train_img/'  #保存train_img
save_train_label = 'car/save_train_label/'#保存train_label
save_test_img = 'car/save_test_img/' #保存test_img
save_test_label = 'car/save_test_label/'#保存test_label
    
train_percent = 0.7

def mkdir():
    if not os.path.exists(img_path):
        os.makedirs(img_path)
    if not os.path.exists(label_path):
        os.makedirs(label_path)
    if not os.path.exists(save_train_img):
        os.makedirs(save_train_img) 
    if not os.path.exists(save_train_label):
        os.makedirs(save_train_label) 
    if not os.path.exists(save_test_img):
        os.makedirs(save_test_img)
    if not os.path.exists(save_test_label):
        os.makedirs(save_test_label)
    

def main():
    mkdir()
    total_txt = os.listdir(label_path)  #获取txt列表
    num_all_txt = len(total_txt)  #txt数目
    list_all_txt = range(num_all_txt) #range[0,num_all_txt]  
    
    num_train = int(num_all_txt*train_percent)  #tain数据的数目
    train = random.sample(list_all_txt,num_train) #打乱随机抽取num_train数量的数据
    test = [i for i in list_all_txt if i not in train]  #选取test数据
    
    print("训练集数目：{}, 测试集数目：{}".format(len(train),len(test)))
    #normal
    for i in list_all_txt:
        name = total_txt[i][:-4]
       
        if i in train:
            shutil.copyfile(img_path + '/'+ name + '.jpg', save_train_img + '/'+ name + '.jpg')
            shutil.copyfile(label_path +'/' + name + '.txt',save_train_label + name +'.txt')
        else:
            shutil.copyfile(img_path + '/'+ name + '.jpg',save_test_img + name + '.jpg')
            shutil.copyfile(label_path +'/' + name + '.txt',save_test_label + name + '.txt')
        
if __name__=="__main__":
    main()