动手学深度学习（pytorch土堆）-01Dataset类代码

ptyorch加载数据

from torch.utils.data import Dataset
# help(Dataset)
import cv2
import os
from PIL import Image


"""
这是类的构造函数，__init__ 方法用于初始化类的对象。
root_dir: 根目录，存储数据集的路径。
label_dir: 标签目录（比如分类标签，如 "ants"）。
self.path: 将 root_dir 和 label_dir 拼接成一个完整的路径。
self.image_path: 使用 os.listdir() 获取路径下所有文件的文件名列表。
"""
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.image_path=os.listdir(self.path)

    def __getitem__(self, idx):
        img_name=self.image_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.image_path)

"""
__getitem__ 是 Dataset 类的核心方法，用于根据索引 idx 获取数据集中的一个样本。
img_name: 获取第 idx 个图像的文件名。
img_item_path: 生成该图像的完整路径。
img: 使用 PIL.Image.open() 方法打开图像。
label: 返回图像所属的标签（这里直接使用文件夹名 label_dir 作为标签）。
return img, label: 返回图像和对应的标签。
"""

root_dir="hymenoptera_data/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)
all_dataset=ants_dataset + bees_dataset#数据集拼接
img1,label1=all_dataset[5]
img2,label2=all_dataset[125]
img1.show()
img2.show()