pytorch的Dataset的使用

from torch.utils.data import Dataset
from PIL import Image
import os
class MyData():
    def __init__(self, root_dir, label_dir):
        #使用self相当于创建一个全局变量给后面的也可以使用
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.img_path = os.listdir(self.path)
        #获取其中每一个图片
    def __getitem__(self, index):
        #通过index获取图片地址
        img_name = self.img_path[index] #每个图片的路径
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        print(img_item_path)
        #读取图片
        img = Image.open(img_item_path)
        img.show
        label = self.label_dir
        return img, label
    def __len__(self): #数据集有多长
        return len(self.img_path)

root_dir = "D:\\数据集\\hymenoptera_data\\train"
ant_lable_dair = "ants"
bees_label_dir = "bees"
ant_dataset = MyData(root_dir, ant_lable_dair)
bees_dataset = MyData(root_dir,bees_label_dir)
ant_dataset.__getitem__(0)
for i in range(1,100):
    bees_dataset.__getitem__(i)
img, label = ant_dataset[0]
img.show()