Как разделить набор данных изображения в X_train, y_train, X_test, y_test?

Question

У меня есть набор данных, подобный следующей структуре:

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)

Как я могу это сделать?

Answer 1

Вы можете построить массивы X и y, используя модуль os:

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 класс для создания набора данных.

Answer 2

Вы всегда можете использовать scikit-learn: https://scikit -learn.org / stable / modules / generate / sklearn.model_selection.train_test_split.html

не забудьте импортироватьэто

from sklearn.model_selection import train_test_split