Dataset类

from torch.utils.data import Dataset
from PIL import Image
import os

#创建MyData类继承Dataset类
class MyData(Dataset):
    #实例化对象时的初始化函数(构造方法)
    def __init__(self, dataset_root_path, dataset_child_path):
        self.dataset_root_path = dataset_root_path
        self.dataset_child_path = dataset_child_path
        #join用于字符串拼接
        self.dataset_path = os.path.join(dataset_root_path,dataset_child_path)
        self.img_path_list = os.listdir(self.dataset_path)

    #获取单个样本数据的方法，index为样本索引
    def __getitem__(self, index):
        img_path = os.path.join(self.dataset_root_path,self.dataset_child_path,self.img_path_list[index])
        img = Image.open(img_path)
        label = self.dataset_child_path
        return img, label

    def __len__(self):
        return len(self.img_path_list)

#实例化对象
ants_dataset = MyData("dataset/train","ants_image")
bees_dataset = MyData("dataset/train","bees_image")