Dataset 读取数据

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

class Mydata(Dataset):
    def __init__(self,root_dir,label_dir):
        self.root_dir = root_dir      #根目录 dataset/train
        self.label_dir = label_dir      #标签的后面链接目录 ants_image
        self.path = os.path.join(self.root_dir,self.label_dir)      #根目录与标签目录链接起来 dataset/train + ants_image
        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)  #该图片的路径 dataset/train/ants_image/5650366_e22b7e1065.jpg
        img = Image.open(img_item_path)   #打开该图片
        label = self.label_dir           #此图片属于哪一个标签 ants_image
        return img,label

    def __len__(self):
        return len(self.img_path)    #图片长度

root_dir = "dataset/train"
ants_label_dir ="ants_image"
bees_label_dir ="bees_image"

ants_dataset = Mydata(root_dir,ants_label_dir)
bees_dataset = Mydata(root_dir,bees_label_dir)

train_dataset = ants_dataset + bees_dataset

img,label = train_dataset[110]
img.show()
print(label)

数据集文件夹对应位置，图片是随便找的