pytorch加载数据

# 1.pytorch加载数据Dataset、Dataloader
# Dataset提供一种方式去获取每个数据及其对应的label，告诉我们一共有多少条数据
# Dataloader为后面的网络提供不同的数据形式，它将一批一批数据进行一个打包
import os

from PIL import Image
# 2.数据集常见的两种形式
# 第一种形式：文件夹的名称是它的label
# 第二种形式：label为文本格式、文本名称为图片名称，文本内容为对应的label
from torch.utils.data import Dataset

help(Dataset)

# 3.路径直接加载数据
img_path = 'D:/05深度学习笔记/demo.png'
img = Image.open(img_path)
img.show()


# 4.Dataset加载数据
class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir #self.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="D:/05深度学习笔记/Data/FirstTypeData/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))
print(len(bees_dataset))
train_dataset=ants_dataset+bees_dataset
print(len(train_dataset))

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