File size: 2,076 Bytes
b83d9ec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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)
|