[python]随机按照比例划分图像数据集为训练集和测试集

原创

是念 2023-04-01 07:29:22 ©著作权

文章标签 python 深度学习机器学习数据集 Image 文章分类 JavaScript 前端开发

©著作权归作者所有：来自51CTO博客作者是念的原创作品，请联系作者获取转载授权，否则将追究法律责任

最近有人跟我发了一个图像数据集，但是没有划分成训练集和测试集，以AnimalFace.zip为示例，数据集的目录为：

....
├── TigerHead
│   ├── 1235_white-baby-tiger.jpg
│   ├── _1768316_tiger4.jpg
│   ├── 2_tigers.jpg
│   ├── bengal_tiger_31.jpg
│   ├── Bengal_Tiger.jpg
.....

我的根目录是Images，其中TigerHead是类别名，下面的jpg是该目录的图片。python代码实现为：

import os
import random
import shutil
import glob
import csv
import numpy as np
def CopyFile(imageDir,test_rate,save_test_dir,save_train_dir):
    #三个参数，第一个为每个类别的所有图像在计算机中的位置
    #第二个为copy的图片数目所占总的比例，最后一个为移动的图片保存的位置，
    image_number = len(imageDir)  #图片总数目
    test_number = int(image_number * test_rate)#要移动的图片数目
    print("要移动到%s目录下的图片数目为:%d"%(save_test_dir,test_number))
    random.seed(42)
    test_samples = random.sample(imageDir, test_number)#随机截取列表imageDir中数目为test_number的元素
    # copy图像到目标文件夹
    if not os.path.exists(save_test_dir):
        os.makedirs(save_test_dir)
        print("save_test_dir has been created successfully!")
    else:
        print("save_test_dir already exited!")
    if not os.path.exists(save_train_dir):
        os.makedirs(save_train_dir)
        print("save_train_dir has been created successfully!")
    else:
        print("save_train_dir already exited!")
    for i,j in enumerate(test_samples):
        shutil.copy(test_samples[i], save_test_dir+test_samples[i].split("/")[-1])
    print("tests 移动完成！number of examples is {}".format(test_number))
    for train_imgs in imageDir:
        if train_imgs not in test_samples:
            shutil.copy(train_imgs, save_train_dir+train_imgs.split("/")[-1])
    print("train 移动完成! number of examples is {}".format(image_number-test_number))

file_path="Image/*/*.jpg"
test_rate = 0.2
file_names = glob.glob(file_path)
save_test_dir='dataset/test/'
save_train_dir='dataset/train/'
CopyFile(file_names,test_rate,save_test_dir,save_train_dir)