如何对自定义图像数据集进行训练测试拆分及独热编码？

1. 提取自变量（X）与因变量（y） #

首先明确两个核心变量：

• 自变量X：所有图像的像素数据，最终是形状为(样本总数, 图像高度, 图像宽度, 通道数)的numpy数组（RGB图通道数为3）
• 因变量y：对应每个图像的类别标签，先给每个类别分配整数ID（比如A→0，B→1，C→2，D→3）

如果你的数据集还没整理成数组格式，用这段代码手动遍历文件夹生成：

import os
from PIL import Image
import numpy as np

dataset_path = "path/"
categories = ["A", "B", "C", "D"]

X = []
y = []

# 遍历每个类别文件夹
for label_idx, category in enumerate(categories):
    class_folder = os.path.join(dataset_path, category)
    # 遍历文件夹内所有图像
    for img_filename in os.listdir(class_folder):
        img_path = os.path.join(class_folder, img_filename)
        # 打开图像并统一尺寸（可根据需求修改尺寸）
        img = Image.open(img_path).resize((224, 224))
        # 转成numpy数组加入X
        X.append(np.array(img))
        # 加入对应类别标签
        y.append(label_idx)

# 转成numpy数组，方便后续操作
X = np.array(X)
y = np.array(y)

2. 拆分训练集与测试集 #

用sklearn的train_test_split快速拆分，还能保证训练/测试集的类别分布一致：

from sklearn.model_selection import train_test_split

# test_size=0.2 表示测试集占20%，random_state固定随机种子保证结果可复现
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

3. 独热编码类别标签 #

独热编码会把整数标签（如0）转成二进制向量（如[1,0,0,0]），适合多数分类模型的输入需求，给你两种常用方法：

方法1：用sklearn的OneHotEncoder（通用机器学习场景） #

from sklearn.preprocessing import OneHotEncoder

# 初始化编码器，sparse_output=False输出密集数组
encoder = OneHotEncoder(categories="auto", sparse_output=False)
# 训练编码器并转换训练集标签
y_train_onehot = encoder.fit_transform(y_train.reshape(-1, 1))
# 转换测试集标签（注意只用transform，不要fit）
y_test_onehot = encoder.transform(y_test.reshape(-1, 1))

# 查看编码器对应的类别映射
print("类别映射：", dict(zip(categories, encoder.categories_[0])))

方法2：用Keras的to_categorical（深度学习场景） #

如果用TensorFlow/Keras做深度学习，直接用这个更方便：

from tensorflow.keras.utils import to_categorical

# num_classes=4对应你的4个类别
y_train_onehot = to_categorical(y_train, num_classes=4)
y_test_onehot = to_categorical(y_test, num_classes=4)

内容的提问来源于stack exchange，提问作者Hinnata