认识Dataset类
利用Jupter可以看到官方给出的定义:
DataSet就是用来封装数据的类,主要用来对数据进行相关的自定义操作(比如图片的裁剪、标签的定义等),通过__getitem__函数返回所需要的数据。
这里有一个getitem函数,getitem函数接收一个index,然后返回图片数据和标签,这个index通常是指一个list的index,这个list的每个元素就包含了图片数据的路径和标签信息。
下面进入实战,对蚂蚁和蜜蜂的图片进行图片的读取
导入相对应的包:
import os
from torch.utils.data import Dataset
from PIL import Image
创建Mydata类来继承Dataset类:
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(root_dir,label_dir)
self.img_path = os.listdir(self.path)
def __getitem__(self, item):
img_game = self.img_path[item]
img_item_path = os.path.join(self.root_dir,self.label_dir,img_game)
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'
ants_dataset=Mydata(root_dir,ants_label_dir)
img1,label1=ants_dataset[0]
img1.show()
运行结果:
蜜蜂:
root_dir='dataset/train'
bees_label_dir='bees'
bees_dataset=Mydata(root_dir,bees_label_dir)
img2,label2=bees_dataset[0]
img2.show()
注意:数据集需要先添加到该py文件的文件夹下
运行结果