Pytorch Dataset读取图片数据

本文为看了小土堆的pytorch的教学视频后的个人总结。视频链接。
本文用到的数据集下载地址： 数据下载地址 密码：5suq

• 首先我们先导入一些之后操作需要的模块。

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):
        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)

os.path.join 对路径进行拼接。
os.listdir() 返回指定的文件夹包含的文件名字的列表。

root_dir = "./data/train"
label_dir = "ants"
ants_dataset = MyData(root_dir, label_dir)

img, label = ants_dataset[1]
print(label)
img.show()

这样操作之后便可以打开相应文件夹中的图片。