如何在 X_train、y_train、X_test、y_test 中拆分图像数据集？

【问题标题】：How to split an image dataset in X_train, y_train, X_test, y_test?如何在 X_train、y_train、X_test、y_test 中拆分图像数据集？
【发布时间】：2019-06-17 09:28:58
【问题描述】：

我有一个如下结构的数据集：

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

我想将此数据集拆分为X_train, y_train, X_test, y_test。

这样：

### data: shuffled and split between train and test
(X_train, y_train), (X_test, y_test) = mnist.load_data()

或者，

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

我该怎么做？

【问题讨论】：

标签： python pandas dataset

【解决方案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 或scikit-image 或类似的库）打开图像。

如果您打算使用pytorch 来训练神经网络，您可以使用他们的ImageFolder class 来创建您的数据集。

【讨论】：