python - 将 train_test_split 与本地目录中的图像一起使用

Question

我已从本地目录读取图像，如下所示:

from PIL import Image
import os

root = '/Users/xyz/Desktop/data'

for path, subdirs, files in os.walk(root):
    for name in files:
        img_path = os.path.join(path,name)

我有两个子目录:category-1 和 category-2，每个子目录都包含属于每个类别的图像文件 (.jpg)。

如何将这些图像和两个类别与 train_test_split() 一起使用Scikit-Learn 中的函数？也就是说，要安排训练和测试数据？

谢谢。

Answer 1

您必须从图像中读取像素数据并将其存储在 Pandas DataFrame 或 numpy 数组中。同时，您必须将相应的类别值 category-1 (1) 和 category-2 (2) 存储在列表或 numpy 数组中。这是一个粗略的草图:我假设您有一些商店类别，它们根据图像名称返回1或2。

X = numpy.array([])
y = list()

for path, subdirs, files in os.walk(root):
  for name in files:
    img_path = os.path.join(path,name)
    correct_cat = categories[img_path]
    img_pixels = list(Image.open(img_path).getdata())
    X = numpy.vstack((X, img_pixels))
    y.append(correct_cat)

您正在有效地存储图像像素和类别值(转换为整数)。可能有其他方法可以做到这一点:Check this例如。

一旦有了 X 和 y 列表，您就可以对它们调用 train_test_split

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