【深度学习动手做】Dataset类，读取数据库中带标签的图片

用于将带标签的数据读取，比如我有一个分类“蚂蚁”和“蜜蜂”的任务，其工作目录下有以下：

data文件夹、read_data.py ...

其中data文件夹下一级又分为“train”，"val”，然后按照标签分为两个文件夹“ants”，“bees”，里面各自是很多张蚂蚁/蜜蜂的图片。

以下代码时read_data.py的内容，用于读取特定文件夹的图片数据。

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   #train或val
        self.label_dir = label_dir   #ant或bee
        self.path = os.path.join(self.root_dir,self.label_dir)
        self.img_path = os.listdir(self.path ) #把path下的所有图片的文件名形成一个list

    def __getitem__(self, idx):
        img_name = self.img_path[idx] #idx就是第几张图片，即index
        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)

 例子：获取训练集中的ant数据集图片

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

得到其中一个实例可以这么做：

为什么这里直接切片取值 bees_dataset[1] 就可以得到实例了呢？不应该使用 __getitem__()函数吗？

原来，__getitem__的是类中一个特殊方法，可以像字典对象那样根据key取值（dict['key']），如类对象Object['key']，系统会自动调用__getitem__方法(所以pytorch里面Dataloader有时候没有发现主函数直接调用__getitem__，原来是系统自动调用！)，然后返回该方法定义return值。