Pytorch土堆学习笔记——P6Dataset类代码实战

关于class，实例化，init，getitem等

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
        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,idx):
    #__getitem__肯定要带idx，或者用key、item等都行
    #一般如果想使用索引访问元素时，就可以在类中定义这个方法（__getitem__(self, key) ）
    #通过键来获取对应的值，idx是键。
    # 当对象是序列时，键是整数。当对象是映射时（字典），键是任意值。
       #idx=1 #这里指定的话，等下在后面ants_dataset[几]都是同样一个图片了
        img_name = self.img_path[idx]
        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 = "train/"
ants_label_dir = "ants"
ants_dataset = MyData(root_dir,ants_label_dir)
#dataset = MyData() 会报错，x=类名(init函数的参数)，x是对类的实例化
img, label = ants_dataset[0]
#当对实例对象进行[idx]操作时，就会调用__getitem__
print(img,'\n',label)
img.show()

print(ants_dataset.__len__())