如何为CNN构建自定义图像数据集并划分训练测试集

一、先理清楚数据集标准结构 #

先把你的两类图像整理成框架能识别的标准分类结构，这是后续操作的基础：

• 新建根文件夹，比如命名为custom_dataset
• 在根文件夹下创建train（训练集）和test（测试集）两个子文件夹
• 分别在train和test内部再建两个子文件夹：own_images（存自有图像）、random_images（存随机图像）
• 后续把两类图像按比例分配到对应的子文件夹里即可

二、训练/测试集的两种划分方式 #

1. 手动划分（适合小数据集） #

• 按7:3或8:2的比例拆分：比如自有图像有100张，就挑70张放到train/own_images，剩下30张放到test/own_images；随机图像做同样操作
• 注意：拆分时要保证两类图像在训练集、测试集中的占比尽量一致，避免数据倾斜影响模型效果

2. 代码自动划分（适合大数据集） #

用Python脚本自动完成，避免手动操作的误差：

import os
import shutil
from sklearn.model_selection import train_test_split

# 替换成你自己的原始图像文件夹路径
original_own_dir = "你的自有图像文件夹路径"
original_random_dir = "你的随机图像文件夹路径"

# 目标数据集根路径
target_root = "custom_dataset"
# 构建各子文件夹路径
train_own = os.path.join(target_root, "train", "own_images")
test_own = os.path.join(target_root, "test", "own_images")
train_random = os.path.join(target_root, "train", "random_images")
test_random = os.path.join(target_root, "test", "random_images")

# 创建所有需要的文件夹（已存在则跳过）
for dir_path in [train_own, test_own, train_random, test_random]:
    os.makedirs(dir_path, exist_ok=True)

# 拆分自有图像
own_img_list = [f for f in os.listdir(original_own_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
train_own_imgs, test_own_imgs = train_test_split(own_img_list, test_size=0.3, random_state=42)
# 复制文件到对应文件夹
for img in train_own_imgs:
    shutil.copy(os.path.join(original_own_dir, img), train_own)
for img in test_own_imgs:
    shutil.copy(os.path.join(original_own_dir, img), test_own)

# 拆分随机图像（逻辑同上）
random_img_list = [f for f in os.listdir(original_random_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
train_random_imgs, test_random_imgs = train_test_split(random_img_list, test_size=0.3, random_state=42)
for img in train_random_imgs:
    shutil.copy(os.path.join(original_random_dir, img), train_random)
for img in test_random_imgs:
    shutil.copy(os.path.join(original_random_dir, img), test_random)

• 说明：test_size=0.3表示测试集占30%，random_state=42是固定随机种子，保证每次拆分结果一致；记得替换代码里的路径为你自己的实际路径

三、用CNN加载划分好的数据集（以TensorFlow/Keras为例） #

结构整理好后，就可以用深度学习框架加载数据集了：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 训练集数据增强（可选，提升模型泛化能力）
train_datagen = ImageDataGenerator(
    rescale=1./255,  # 像素值归一化到0-1
    rotation_range=20,  # 随机旋转±20度
    width_shift_range=0.2,  # 随机水平偏移20%
    height_shift_range=0.2,  # 随机垂直偏移20%
    horizontal_flip=True  # 随机水平翻转
)

# 测试集只做归一化，不做数据增强
test_datagen = ImageDataGenerator(rescale=1./255)

# 加载训练集
train_loader = train_datagen.flow_from_directory(
    os.path.join(target_root, "train"),
    target_size=(150, 150),  # 统一图像尺寸，根据你的模型输入调整
    batch_size=32,
    class_mode='binary'  # 二分类问题用binary，多分类用categorical
)

# 加载测试集
test_loader = test_datagen.flow_from_directory(
    os.path.join(target_root, "test"),
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary'
)

四、常见失败原因排查 #

你之前划分失败可能是这些问题：

• 原始文件夹里混了非图像文件，代码没过滤导致报错——解决：添加文件后缀判断（如代码里的.jpg/.png过滤）
• 路径写错，程序找不到文件——解决：打印路径检查是否正确
• 两类数据量差距过大，导致模型偏向多数类——解决：尽量保证两类数据量均衡，或给少数类加损失权重

内容的提问来源于stack exchange，提问作者Pothula Vinitha reddy