📌  相关文章
📜  如何将图像数据集拆分为训练集和测试集 keras - Python (1)

📅  最后修改于: 2023-12-03 15:24:44.606000             🧑  作者: Mango

如何将图像数据集拆分为训练集和测试集 keras - Python

在机器学习中,我们需要使用一个已经处理好的数据集来训练模型以便让其对未知数据进行预测。通常,我们需要随机将数据集分为训练集和测试集。

在本篇文章中,我们将会学习如何使用 keras 和 Python 中的 scikit-learn 库,将图像数据集拆分为训练集和测试集。

准备工作

首先,我们需要导入模块,并将数据集加载到程序中。

import os
import cv2
import random
import numpy as np
from keras.utils import np_utils
from sklearn.model_selection import train_test_split

path = 'path/to/dataset'
categories = os.listdir(path)
labels = [i for i in range(len(categories))]

label_dict = dict(zip(categories, labels))

img_size = 100
data = []
target = []

for category in categories:
    folder_path = os.path.join(path, category)
    img_names = os.listdir(folder_path)

    for img_name in img_names:
        img_path = os.path.join(folder_path, img_name)
        img = cv2.imread(img_path)

        try:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            resized = cv2.resize(gray, (img_size, img_size))
            data.append(resized)
            target.append(label_dict[category])
        except Exception as e:
            pass

data = np.array(data) / 255.0
data = np.reshape(data, (data.shape[0], img_size, img_size, 1))
target = np.array(target)
target = np_utils.to_categorical(target)

在这里,我们首先定义了数据集的路径 path,数据集中各分类的文件夹 categories,以及每个分类的标签值 labels。接着,我们将分类和标签值打包成字典 label_dict

然后,我们将读取所有图像,并将其转换为灰度图像。接着,我们将图像大小重新调整为指定的大小,并将其添加到数据列表 data 和标签列表 target 中。

最后,我们将数据列表和标签列表转换为 numpy 数组,并将像素值缩放到 0 到 1 之间。我们还将标签转换为 one-hot 编码形式。

将数据集拆分为训练集和测试集

接下来,我们使用 scikit-learn 库中的 train_test_split 函数,将数据集拆分为训练集和测试集。

train_data, test_data, train_target, test_target = train_test_split(data, target, test_size=0.2, random_state=42)

在这里,我们将数据集和标签集传递给 train_test_split 函数,并将测试集占总数据的 20%。我们还指定了一个随机种子 random_state,以确保每次拆分的结果一致。

总结

这就是将图像数据集拆分为训练集和测试集的完整指南。您可以将上述代码添加到您自己的项目中,并根据自己的需求进行自定义。

需要注意的是,训练数据集的大小对模型的性能和泛化能力有很大的影响。因此,请确保您的训练集足够大,以避免过拟合和欠拟合的问题。