创建训练和测试的数据集，并创建加载器。lib.load_imags.py:

import glob
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from script.setting import *


# 图片处理函数
def img_loader(path):
    try:
        img = Image.open(path)
        img = img.convert('RGB')  # 转换成RGB模式
        if img.size != normalize_size:  # 标准化图片尺寸
            img = transforms.Resize(normalize_size)(img)  # 提供新尺寸
        return img
    except Exception as e:
        print(f"Error loading image {path}: {e}")
        return error_default_img  # 返回一个默认的错误图像


# 训练集数据预处理方法
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.RandomVerticalFlip(),  # 随机垂直翻转
    transforms.RandomRotation(90),  # 随机旋转90度
    # transforms.RandomGrayscale(p=0.1),  # 随机将图片转换为灰度图，p=0.1表示有10%的概率执行该操作
    # transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5),  # 调整图片的亮度、对比度、饱和度和色调
    transforms.ToTensor(),  # 将图片转换为Tensor
    transforms.Normalize(normalize_mean,  # 标准化
                         normalize_std)
])

# 测试集数据预处理方法
test_transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片转换为Tensor
    transforms.Normalize(normalize_mean,  # 标准化
                         normalize_std)
])


# 自定义数据集
class MyDataset(Dataset):
    def __init__(self, img_list,  # 图片的地址列表
                 transform=None):
        super(MyDataset, self).__init__()

        imgs = []  # 图片的路径和类别标签列表（转换成了0-9的数字的）存储格式为：
        # [[图片路径, 0-9的图片类别], ...]

        for img_path in img_list:
            im_label_name = img_path.split('\\')[-2]  # 图片所属类别的名称，这里使用的是绝对路径，
            # 文件目录分隔符为反斜杠，使用相对路径则为正斜杠
            imgs.append([img_path, label_dict[im_label_name]])  # 将图片路径和对应的类别标签添加到列表中

        self.imgs = imgs
        # print(imgs[0])
        self.transform = transform
        self.loader = img_loader  # 图片加载函数

    def __getitem__(self, index):  # 获取数据集的图片和标签
        img_path, label = self.imgs[index]
        img_data = self.loader(img_path)  # 图片的数据

        if img_data is None:
            return None, label  # 处理加载错误的情况

        if self.transform is not None:
            img_data = self.transform(img_data)  # 图片数据转换

        return img_data, label  # 返回图片数据和标签

    def __len__(self):  # 获取数据集的图片数量  train_path
        return len(self.imgs)


# 获取训练集的文件名
train_list = glob.glob(train_path + '\\*\\*')

# 获取测试集的文件名
test_list = glob.glob(test_path + '\\*\\*')

# 定义训练数据集
trans_dataSet = MyDataset(train_list, transform=train_transform)
# 训练集图片数量
train_num = len(trans_dataSet)

# 定义测试数据集
test_dataSet = MyDataset(test_list, transform=test_transform)
# 测试集图片数量
test_num = len(test_dataSet)

# 定义训练集的加载器
train_loader = DataLoader(trans_dataSet, batch_size=batch_size, shuffle=True, num_workers=num_workers)

# 定义测试集的加载器
test_loader = DataLoader(test_dataSet, batch_size=batch_size, shuffle=False, num_workers=num_workers)

# print("num_of_train", len(train_loader))  # 391(50000/128)，相当于有391个batch，每个batch有128个样本
# print("num_of_test", len(test_loader))  # 79(10000/128)，相当于有79个batch，每个batch有128个样本