【Pytorch】加载数据集 - 学习笔记

大佬的原视频
之前学习的代码里面，数据集已经都加载好了，不需要自己动手写，所以一直想学习怎么加载自己的数据集，记录一下学习的视频，免得之后忘了

准备：数据集格式

蚂蚁和蜜蜂的分类，原始数据集仅仅包含图片文件夹(ants_image和bees_image)，标签文件夹是后来加的，可以先不看

图片的文件夹显示如下

读取数据集

1、把数据集加载到项目中

点击Explorer，找到当前项目的文件夹，然后把Dataset文件夹丢进去

2、写代码加载数据集

from torch.utils.data import Dataset
import os  #python关于系统的库
from PIL import Image

class MyData(Dataset):
    # 初始化类 根据类创建实例时要运行函数，为整个class提供全局变量
    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir  # 函数的变量不能传递给另外一个变量，而self能够把指定变量给别的函数使用，全局变量
        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] #获取图片名称  self.全局的
        img_item_path = os.path.join(self.path, img_name) # 获取每个图片的地址(相对路径)
        img = Image.open(img_item_path)
        label = self.label_dirA
        return img, label

    def __len__(self):
        return len(self.img_path)

root_dir = "dataset/hymenoptera_data/train"
ants_label_dir = "ants_image"
bees_label_dir = "bees_image"
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[123]
img.show()

输出结果为

【附】3、如果想把label单独放一个文件夹呢

目标：在只有图片数据集的情况下，形成下面这样的label文件


只需要运行下面的代码即可

import os

root_dir = "dataset/hymenoptera_data/train"
target_dir = "ants_image"
img_path = os.listdir(os.path.join(root_dir, target_dir))
label = target_dir.split('_')[0]
out_dir = "ants_label"

for i in img_path:
    file_name = i.split('.jpg')[0]
    with open(os.path.join(root_dir, out_dir, "{}.txt".format(file_name)), 'w') as f:
        f.write(label)