04 Dataset类代码实战

创建一个类MyData，这个类继承自Dataset类，并创建一个初始化函数，相当于C++中的构造函数。

class MyData(Dataset):

	det __init__(self):

	def __getitem__(self, idx):

然后定义一个getitem函数，idx相当于索引。

为了获取图片地址的列表，我们需要使用os工具包，在最前面加上

import os

然后就可以在初始化函数中通过图片路径和标签路径来对图片列表进行导入了，主要导入为类中的全局变量，类似于C++中的属性(self.后面)

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) //获取图片地址列表

这样就可以在getitem函数中通过idx索引获取具体某一张图片的信息

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

最后定义一个len函数，用于获取数据长度

def __len__(self):
	return len(self.img_path)

我们可以尝试获取图片信息

root_dir = "dataset/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

需要注意的是，这种图片获取方式针对的是图片的标签直接用文件夹命名的情况

但现实中我们更多的会将标签用一个文档存储起来一一对应，比如这样

此时label文件夹中的文档保存的就是相应图片的标签，文档名等于相应图片名

我们可以用如下代码将标签存入文档中

import os  
  
root_dir = "dataset/train"  
  
target_dir  = "ants_image"  
  
img_path = os.listdir(os.path.join(root_dir, target_dir))  
label = target_dir.split('_')[0]  
out_dir = "ants_label"  
for i in img_path:  
    file_name = i.split('.jpg')[0]  
    with open(os.path.join(root_dir, out_dir, "{}.txt".format(file_name)), 'w') as f:  
        f.write(label)