六、Dataset类实战

用Dataset类读取图片：

数据集地址：链接：https://pan.baidu.com/s/1R5mBwBo-AGydrwZX0EDbCA     提取码：tkrp

 

from torch.utils.data import Dataset        #加载Dataset类，在torch工具箱内，常用的工具区utils,其中关于数据的data区中导入
import os                                   #打开文件夹。获取所有图片的地址，python中关于系统的一个库
from PIL import Image                       #读取图片。或者用opencv,写为import cv2
class Mydata(Dataset):                      #继承Dataset类
    def __init__(self,root_dir,lable_dir):  #类的实例化，主函数创建实例对象时，自动调用该函数，主要作为外部信息传入类中私有方法的接口
        self.root_dir=root_dir              #图片的目录地址。函数中的变量不能赋值给另一个函数，但self可以
        self.lable_dir=lable_dir            #图片的标签信息
        self.img_name=os.listdir(self.root_dir)#os打开目录，listdir函数返回目录下文件（夹）的名称列表
    def __getitem__(self, idx):             #实例化__getitem__方法，idx代表索引。实力对象通过：对象[idx]，python会自动调用该函数
        img_name=self.img_name[idx]         #列表中索引为idx的值（文件名称，此处值图片名称）赋给img_name
        img_path=os.path.join(self.root_dir,img_name)#具体的图片路径，用os.path.join方法连接两个路径
        img=Image.open(img_path)            #打开图片，将图片信息赋值给img,具体图片包含的信息可以在python console中查看
        lable = self.lable_dir              #lable赋值
        return img,lable                    #返回图片信息和标签
    def __len__(self):                      #实例化__len__方法
        return len(self.img_name)           #返回一共多少张图片，即列表的长度
ants_root_dir="E:\\pytorch pycharm\\hymenoptera_data\\train\\ants" #图片文件目录信息。一般图片复制到project下可以复制相对路径
ants_lable_dir="ants"                            #索引信息
ants_dataset=Mydata(ants_root_dir,ants_lable_dir)     #创建蚂蚁的实例对象
img,lable=ants_dataset[0]                   #ants_dataset[idx]自动调用__getitem__方法
img.show()                                  #img.show()方法展示图片
bees_root_dir="E:\\pytorch pycharm\\hymenoptera_data\\train\\bees"
bees_lable_dir="bees"
bees_dataset=Mydata(bees_root_dir,bees_lable_dir)     #创建蜜蜂的实例对象
train_dataset=ants_dataset+bees_dataset   #+号可以把两个数据集合并，索引范围变为两类图片数量之和，通常用于训练时图片数量不足，补充仿造的数据集的情况
img,lable=train_dataset[124]              #索引号增加
img.show()