Pytorch制作数据集

from torch.utils.data import Dataset#import Dataset类
from PIL import Image#图像处理
import os#操作系统相关库，用来根据路径读取数据

class Mydata(Dataset):#我们读取数据的类要继承Dataset类
    def __init__(self,root_dir,label_dir):
        self.root_dir=root_dir
        self.label_dir=label_dir
        self.path=os.path.join(root_dir,label_dir)#将两个路径合并为一个
        self.img_path=os.listdir(self.path)#获取所有图片的文件名列表

    def __getitem__(self, idx):#返回数据标签对的函数
        img_name=self.img_path[idx]#通过idx(索引)访问数据，最终实例化Mydata后可直接##通过索引访问数据-标签对
        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)#蜜蜂数据集
train_dataset=ants_dataset+bees_dataset#合并两个数据集
img,label=ants_dataset[0]#通过索引读取数据对
img.show()#打印图片
print(label)#打印label