深度学习Dataset类

 蚂蚁蜜蜂分类数据集和下载连接https://download.pytorch.org/tutorial/hymenoptera_data.zip

from torch.utils.data import Dataset
from PIL import Image
import os
# 继承Dataset
class MyData(Dataset):
    # 初始化 为整个class提供全局变量
    def __init__(self,root_dir,label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        # 获取图片路径地址 'dataset/train\\ants'
        self.path = os.path.join(self.root_dir,self.label_dir)
        # 将文件夹下所有地址变成一个列表 img_path_list[0] = '0013035.jpg'
        self.img_path = os.listdir(self.path)

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        # 每个图片路径
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        # 读取图片
        img = Image.open(img_item_path)
        label = self.label_dir
        return img,label

    def __len__(self):
        # 列表的长度
        return len(self.img_path)

root_dir = "dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = MyData(root_dir,ants_label_dir)
bees_dataset = MyData(root_dir,bees_label_dir)
# image,label = ants_dataset[0] img.show() 打开图片
train_dataset = ants_dataset + bees_dataset