如何从一个大文件夹中读取图片并将其拆分为训练集、验证集和测试集?

How to read pictures from a big folder and split it into train, validation and test sets?

我正在使用 pytorch 进行手语手势分类,我有类似于每个字母的图片,这些图片位于以该特定字母命名的文件夹中。例如。文件夹 "A" 有“1_A_1.jpg”、“1_A_2.jpg”、“21_A_3.jpg”等

我正在尝试构建一个函数:

  1. 遍历不同的文件夹
  2. 将数据拆分为训练集、验证集和测试集
  3. 用它们各自的文件夹名称标记这些图片(即字母标签)
  4. Returns 3 个创建的文件夹,分别是训练、测试和验证

所有在线代码都显示了拆分来自 torchvision 数据集(内置数据集)的数据的示例,没有从头开始。

我在 Whosebug 上找到了以下内容:

import os
import numpy as np
import argparse

def get_files_from_folder(path):

    files = os.listdir(path)
    return np.asarray(files)

def main(path_to_data, path_to_test_data, train_ratio):
    # get dirs
    _, dirs, _ = next(os.walk(path_to_data))

    # calculates how many train data per class
    data_counter_per_class = np.zeros((len(dirs)))

    for i in range(len(dirs)):
        path = os.path.join(path_to_data, dirs[i])
        files = get_files_from_folder(path)
        data_counter_per_class[i] = len(files)
    test_counter = np.round(data_counter_per_class * (1 - train_ratio))

    # transfers files
    for i in range(len(dirs)):
        path_to_original = os.path.join(path_to_data, dirs[i])
        path_to_save = os.path.join(path_to_test_data, dirs[i])

        #creates dir
        if not os.path.exists(path_to_save):
            os.makedirs(path_to_save)
        files = get_files_from_folder(path_to_original)
        # moves data
        for j in range(int(test_counter[i])):
            dst = os.path.join(path_to_save, files[j])
            src = os.path.join(path_to_original, files[j])
            shutil.move(src, dst)

当我尝试执行以下操作时:

path_to_data= r'path\A'
path_to_test_data=r"path\test"
train_ratio=0.8

main(path_to_data,path_to_test_data,train_ratio)

什么都没发生..

如果我能让它用于训练和测试,我可以轻松地扩展它以进行验证。

试一试:

from pathlib import Path

def main(data_path, out_path, train_ratio):
    #1
    dir_paths = [child for child in Path(data_path).iterdir() if child.is_dir()]

    for i, dir_path in enumerate(dir_paths):
        #2
        files = list(dir_path.iterdir())
        train_len = int(len(files) * (1 - train_ratio)) 

        #3
        out_dir = Path(out_path).joinpath(dir_path.name)
        if not out_dir.exists():
            out_dir.mkdir(parents=True)

        #4
        for file_ in files[:train_len]:
            file_.replace(out_dir.joinpath(file_.name))

if __name__ == '__main__':
    main('data', 'test', 0.8)