tools/将数据集按照比例进行拆分.py

import os
import sys
import numpy as np

from tqdm import tqdm
import cv2
import matplotlib.pyplot as plt

import torchvision
from torchvision import models,transforms,datasets
import torch
import torch.nn as nn
from torch import optim





import os
import  shutil
import random
def make_dir(path):
    import os
    dir = os.path.exists(path)
    if not dir:
        os.makedirs(path)


def get_filename_and_houzhui(full_path):
    import os
    path, file_full_name = os.path.split(full_path)
    file_name, 后缀名 = os.path.splitext(file_full_name)
    return path,file_name,后缀名


dataset_root_path = '../data/cat_vs_dog'
train_path_cat_new = os.path.join(dataset_root_path, 'new/train/cat')
train_path_dog_new = os.path.join(dataset_root_path, 'new/train/dog')

test_path_cat_new = os.path.join(dataset_root_path, 'new/test/cat')
test_path_dog_new = os.path.join(dataset_root_path, 'new/test/dog')

make_dir(train_path_cat_new)
make_dir(train_path_dog_new)
make_dir(test_path_cat_new)
make_dir(test_path_dog_new)

image_dir_path = os.path.join(dataset_root_path,'train')
image_name_list = os.listdir(image_dir_path)



for image_name in tqdm(image_name_list):
    image_path = os.path.join(image_dir_path,image_name)
    path, file_name, 后缀名 = get_filename_and_houzhui(full_path=image_path)
    # print("file_name:", file_name)
    # 定义随机数的范围和对应的概率
    nums = [1, 2]
    probs = [0.9, 0.1]      #设定训练集和测试集的比率

    random_nums = random.choices(nums, weights=probs)[0]

    if(random_nums == 1):       #摇筛子如果摇到了1，那么就是训练集

        if('cat' in file_name):
            shutil.copy(image_path, train_path_cat_new)
        elif('dog' in file_name):
            shutil.copy(image_path, train_path_dog_new)
    elif(random_nums == 2):         #摇骰子如果摇到了2，那么就是测试集
        if('cat' in file_name):
            shutil.copy(image_path, test_path_cat_new)
        elif('dog' in file_name):
            shutil.copy(image_path, test_path_dog_new)