Dataset类实践
蚂蚁蜜蜂分类数据集和下载链接https://download.pytorch.org/tutorial/hymenoptera_data.zip
Dataset:提供一种方式去获取数据及其lable
-
Q:如何获取每个数据及其lable
重写构造方法和获取标签方法
-
Q:告诉我们总共有多少数据
重写len方法
代码示例:
from torch.utils.data import Dataset
from PIL import Image
import osclass MyData(Dataset):def __init__(self, root_dir, label_dir): # 获取所有图片的地址self.root_dir = root_dir # self设置成全局变量self.label_dir = label_dirself.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, ants_label_dir) # 获取蜜蜂的数据集#获取整个train数据集
train_dataset = ants_dataset + bees_dataset
在控制台中进行测试
- 对数据集中图片的相关操作
获取蚂蚁数据集
- 查找当前数据集中第一个图片名称
- 图片名称拼接(进行路径和标签的拼接)
- 读取图片相应信息
实例化对象
- 返回该对象 image和label
- 结果
- 改变ants_dataset[],展示第二张图
获取蜜蜂的数据集
获取整个train数据集(蚂蚁+蜜蜂)
train_dataset = ants_dataset + bees_dataset
# 进行长度测试
len(train_dataset)
248
len(ants_dataset)
124
len(bees_dataset)
124