前言

在训练的时候通常使用Dataset来处理数据集。

Dataset的作用

提供一个方式获取数据内容和标签（label）。

实战

from torch.utils.data import Dataset

from PIL import Image
import os

class get_data(Dataset):

    def __init__(self,root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.img_dir = os.path.join(root_dir,label_dir)
        self.img_list = os.listdir(self.img_dir)


    def __getitem__(self, indx):

        img_path = os.path.join(self.img_dir,self.img_list[indx])
        img_label = self.label_dir
        img_data = Image.open(img_path)
        return img_data,img_label

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

root_dir = "C:\\Users\\Traveler\\Pictures"
label_dir = "Screenshots"

test = get_data(root_dir,label_dir)

img , label = test[1]

# img.show()
print(label)
print(len(test))