如何在X轴、y轴、X轴测试、y轴测试中分割图像数据集？

Question

我有一个类似以下结构的数据集：

dataset/
   |
   |
   -----Pothole/
   |         |
   |         ------ umm001.jpg
   |         |
   |         ------ abd.jpg
   |         |
   |         ------ 
   |         |
   |
   |
   ----Road/
         |
         ------road005.jpg
         |
         ------ummm.jpg
         |
         ------
         |

我想把这个数据集分成X_train, y_train, X_test, y_test。在

因此：

^{pr2}$

或者

(X_train, y_train), (X_test, y_test) = train_test_split(X,y, test_size=0.20)

我该怎么做？在

Answer 1

您可以使用os模块来构建X和y数组：

import os

X = []
y = []
base_dir = '<full path to dataset folder>/'
for f in sorted(os.listdir(base_dir)):
    if os.path.isdir(base_dir+f):
        print(f"{f} is a target class")
        for i in sorted(os.listdir(base_dir+f)):
            print(f"{i} is an input image path")
            X.append(base_dir+f+'/'+i)
            y.append(f)
print(X)
print(y)

然后，您可以使用train_test_split(X,y, test_size=0.20)来获得所需的内容，但请记住，您必须使用其他库（如pillow或{}）或类似库打开这些图像。在

如果您计划使用pytorch来训练神经网络，那么可以使用它们的ImageFolderclass来创建数据集。在

Answer 2

您可以始终使用scikit learn:https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html

别忘了进口

from sklearn.model_selection import train_test_split