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