Dataset类:提供了一种方式去获取数据及其label
- 主要应用:
- 如何获取每一个数据和对应的label
- 如何获取数据的长度
导入
from torch.utils.data import Dataset
os.path.join
对路径进行拼接。os.listdir()
返回指定的文件夹包含的文件名字的列表。
#展示Dataset类的用法:提供了一种获取数据的和对应label的方式
""""
1.如何获取每一个数据及其label
2.获取数据的长度
"""
from torch.utils.data import Dataset #引入对应的package
import os
from PIL import Image #图片操作的package
class MyData(Dataset):#建立一个类,要继承Dataset类
#root_dir是train的相对地址 label_dir是ants
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="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)
train_dataset=ants_dataset+bees_dataset
img,label=ants_dataset[0]
#img.show()
print(len(train_dataset))