学习Pytorch加载数据Dataset类代码实战与练习代码

前言

在学习B站pytorch视频https://www.bilibili.com/video/BV1hE411t7RN?p=7时的实战代码与练习代码记录如下：
视频数据集及相关下载：https://pan.baidu.com/s/16koDDDvAFO7rNoi-uR-A_g 提取码：pahi

1、实战代码（标签即文件名）

代码如下（示例）：

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)

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)
print(len(ants_dataset))
img, label = ants_dataset[0]
img.show()
print(label)
train_dataset = ants_dataset + bees_dataset
print(len(train_dataset))
img, label = train_dataset[123]
img.show()
print(label)
img, label = train_dataset[124]
img.show()
print(label)

2.练习代码（标签另存在txt文件中）

代码如下（示例）：

from torch.utils.data import Dataset
from PIL import Image
import os
'''
Dataset类的作用： 提供一种方式去获取数据及其label
    如何获取一个数据及其label
    告诉我们总共有多少个数据
'''

class MyData(Dataset):
    def __init__(self,root_dir,image_dir,label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.image_dir = image_dir
        self.img_path = os.path.join(self.root_dir,self.image_dir)
        self.label_path = os.path.join(self.root_dir,self.label_dir)
        self.img_path_list = os.listdir(self.img_path)
        self.label_path_list = os.listdir(self.label_path)

    def __getitem__(self, idx):
        img_name = self.img_path_list[idx]
        label_name = os.path.join(self.root_dir, self.label_dir, self.label_path_list[idx])
        with open(label_name, "r") as f:
            label = f.read()
        img_item_path = os.path.join(self.root_dir,self.image_dir,img_name)
        img = Image.open(img_item_path)
        return img, label

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

root_dir = "dataset_txt\\train"
ants_image_dir = "ants_image"
ants_label_dir = "ants_label"
bees_image_dir = "bees_image"
bees_label_dir = "bees_label"

ants_dataset = MyData(root_dir, ants_image_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_image_dir, bees_label_dir)

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

train_dataset = ants_dataset + bees_dataset
print(len(train_dataset))
img, label = train_dataset[123]
img.show()
print(label)

img, label = train_dataset[124]
img.show()
print(label)

总结

这里对文章进行总结：