# # --------------------------
# 1、导入所需三方库
# # --------------------------
# 导入 Python 的 os 模块，用于与操作系统进行交互，如文件和目录操作；用于进行文件和目录相关的操作，如获取文件列表、拼接路径等
# 从 Scikit-learn 库的模型选择模块导入 train_test_split 函数，用于划分数据集为训练集和测试集
# 从 Python Imaging Library (PIL) 导入 Image 模块，用于图像的打开、操作和保存
# 导入 numpy 库并将其别名为 np，numpy 是 Python 中用于科学计算的基础库，提供了多维数组对象和各种数学函数
# # --------------------------import os # 用于与操作系统交互，处理文件路径、目录操作等
from sklearn.model_selection import train_test_split # 用于将数据集拆分训练集和测试集
from PIL import Image # 用于图像的打开、保存、格式转换等基础图像处理
import numpy as np # 用于数值计算，处理图像的数组表示（如像素矩阵）
import cv2 # OpenCV库，用于复杂图像处理（如 resize、滤波、边缘检测等）
import shutil # 用于文件和目录的高级操作（如复制、移动、删除等）# # 不同类别数据分布比例要求：所有类别的样本数量相差≤5%。
# 计算 “类别不平衡系数”：用「最大类别样本数 ÷ 最小类别样本数」，结果＜3 为可接受，＞10 为严重不平衡。
# 看 “样本量是否满足训练需求”：即使比例均匀，若某类别样本量过少（如仅 10-20 张），也会因数据不足导致模型欠拟合，通常建议每个类别至少有 100 张以上样本。# # --------------------------
# 2、数据加载
# # --------------------------
#数据集路径
data_dir = r".\data"
# 初始化三个空列表
all_img = []#用于存储所有图像的文件路径
all_label = []#用于存储所有图像对应的标签
class_names = [] # 用于存储类别文件夹名（用于后续创建新目录）# # --------------------------
# 遍历数据文件夹中的所有子文件夹
# enumerate 函数用于同时获取子文件夹的索引和名称，索引将作为图像所属类别的标签
# # --------------------------
for index, file in enumerate(os.listdir(data_dir)):file_path = os.path.normpath(os.path.join(data_dir, file))# 只处理子目录（跳过根目录下的文件）if not os.path.isdir(file_path):continueelif os.path.isdir(file_path):# 将类别名称添加到class_names列表class_names.append(file)# 使用index作为标签（从0开始）# 遍历当前子文件夹中的所有图像文件for img in os.listdir(file_path):if img.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):# 构建当前图像文件的完整路径img_path = os.path.join(file_path, img)# 将当前图像的完整路径添加到 all_img 列表中all_img.append(img_path)# 将当前子文件夹的索引作为标签添加到 all_label 列表中all_label.append(index)# 标签从0开始，与class_names索引对应# # --------------------------
# 检查加载结果
# # --------------------------
print(f"\n=== 数据加载结果 ===")
print(f"总图片数量: {len(all_img)}")
print(f"总类别数量: {len(set(all_label))}")
print(f"类别列表: {sorted(set(all_label))}")# # --------------------------
# 3、划分数据集
# # --------------------------
# 一般:
# 小规模数据集,训练集:测试集=8:2
# 中等规模数据集,训练集:验证集:测试集=6:2:2
# 大规模数据集,训练集:验证集:测试集=98:1:1
# # --------------------------
'''
# random_state=42 用于保证每次划分的结果一致，方便复现实验结果,
随机数种子控制每次划分训练集和测试集的模式，其取值不变时划分得到的结果一模一样，
其值改变时(可以是任意一个整数)，划分得到的结果不同。
若不设置此参数，则函数会自动选择一种随机模式，得到的结果也就不同。
'''
test_size1=0.2 # test_size1=0.2 表示测试集占整个数据集的 20%
test_size2=0.25 # test_size2=0.25 表示验证集占训练集的 25%, 0.8*0.25=0.2,则测试集占总数据集的20%
random_state=42 #用于保证每次划分的结果一致，方便复现实验结果# 使用 train_test_split 函数,将所有图像划分为训练集train_imgs和测试集test_imgs,所有标签划分为训练集标签train_labels和测试集标签test_labels
train_imgs, test_imgs, train_labels, test_labels = train_test_split(all_img, all_label, test_size=test_size1, random_state=random_state)# 再次使用 train_test_split 函数将训练集进一步划分,分为训练集和验证集
train_imgs, val_imgs, train_labels, val_labels = train_test_split(train_imgs, train_labels, test_size=test_size2, random_state=random_state)# --------------------------
# 4、创建新数据集文件夹并复制图像
# --------------------------
def copy_imgs_to_new_dir(img_paths, labels, save_root, class_names):"""将图像复制到新目录img_paths: 图像路径列表（如train_imgs）labels: 对应图像的标签列表（如train_labels）save_root: 保存根目录（如"./new_data/train"）class_names: 原始类别文件夹名列表（如["cat", "dog"]）"""# 1. 遍历每个图像路径和对应标签for img_path, label in zip(img_paths, labels):# 确保标签在有效范围内if label < 0 or label >= len(class_names):print(f"警告：无效标签 {label}，跳过图像 {img_path}")continue# 获取原始图像的文件名（如"img1.jpg"）img_filename = os.path.basename(img_path)# 确定目标类别文件夹（根据标签映射到原始类别名）target_class_dir = os.path.join(save_root, class_names[label])# 2. 创建目标文件夹（如果不存在）if not os.path.exists(target_class_dir):os.makedirs(target_class_dir) # 递归创建多级目录# 3. 复制图像到目标路径target_img_path = os.path.join(target_class_dir, img_filename)shutil.copy2(img_path, target_img_path) # copy2保留文件元信息（推荐）# 定义新数据集的根目录（可自定义路径）
new_data_root = r".\new_data"# 分别复制训练集、验证集、测试集
print("正在复制训练集图像...")
copy_imgs_to_new_dir(img_paths=train_imgs,labels=train_labels,save_root=os.path.join(new_data_root, "train"), # 目标：./new_data/train/类别名/class_names=class_names
)print("正在复制验证集图像...")
copy_imgs_to_new_dir(img_paths=val_imgs,labels=val_labels,save_root=os.path.join(new_data_root, "val"), # 目标：./new_data/val/类别名/class_names=class_names
)print("正在复制测试集图像...")
copy_imgs_to_new_dir(img_paths=test_imgs,labels=test_labels,save_root=os.path.join(new_data_root, "test"), # 目标：./new_data/test/类别名/class_names=class_names
)print(f"\n所有图像复制完成！新数据集路径：{os.path.abspath(new_data_root)}")# # --------------------------
# 5、定义一个函数，用于将图像文件路径列表转换为图像数组列表,并简单处理
# # --------------------------
def data_sort(images):# 初始化一个空列表，用于存储转换后的图像数组images_list = []# 遍历图像文件路径列表for img_path in images:# 01. 读取图像,使用 Image.open 函数打开图像文件img_src = Image.open(img_path)if img is None:raise ValueError(f"无法读取图像: {img_path}")# 02. 将图像调整为 299x299 的大小，以便后续输入到模型中img_resized = img_src.resize((299, 299))# 03. 转换为RGB格式（InceptionV3训练时使用RGB）img_resized=np.array(img_resized)img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)# 04. 将 images_rgb 添加进列表中images_list.append(img_rgb)return np.array(images_list)# # --------------------------
# 6、调用data_sort函数，处理数据集
# # --------------------------# 调用 data_sort 函数处理训练集的图像文件
train_images = np.array(data_sort(train_imgs),dtype=np.float32)#print(train_images.shape)结果(2693, 299, 299, 3)
# 调用 data_sort 函数处理验证集的图像文件
val_images = np.array(data_sort(val_imgs),dtype=np.float32)
# 调用 data_sort 函数处理测试集的图像文件
test_images = np.array(data_sort(test_imgs),dtype=np.float32)# # --------------------------
# 6、数据类型转换
# # --------------------------
# 将训练集的标签列表转换为 numpy 数组，方便后续处理和输入到模型中
train_labels = np.array(train_labels,dtype=np.float32)
# 将验证集的标签列表转换为 numpy 数组，方便后续处理和输入到模型中
val_labels = np.array(val_labels,dtype=np.float32)