pytorch 的Dataset 使用方法

1. 导入模块：

   • torch.utils.data.Dataset：PyTorch 中用于创建自定义数据集的基类。
   • PIL.Image：Python Imaging Library (PIL) 用于处理图像文件。
   • os：提供了一系列与操作系统交互的函数，如文件路径操作。
   • 使用dataset：https://download.pytorch.org/tutorial/hymenoptera_data.zip
   • conda 里面下载openmv的库

     conda install opencv_python
     

2. 自定义数据集类：

   • __init__ 方法：初始化数据集对象，设置根目录和标签目录，并将它们组合成完整的图像路径。然后，通过 os.listdir 获取所有图像文件名。
   • __getitem__ 方法：根据索引 idx 返回一个样本。首先，它根据索引获取图像文件名，然后构造完整的图像文件路径，并使用 Image.open 打开图像文件。最后，它返回图像和对应的标签。
   • __len__ 方法：返回数据集中样本的数量，即图像文件的数量。

3. 数据集实例化：

   • 代码中定义了两个数据集实例 ants_dataset 和 bees_dataset，分别对应蚂蚁和蜜蜂的训练数据。
   • 第30行的 train_dataset = ants_dataset + bees_dataset 尝试将两个数据集相加。

from torch.utils.data import Dataset
from PIL import Image
import os


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.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 =""%%相对路径
ants_label_dir="ants"
bees_lable_dir ='bees'
ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_lable_dir)

train_dataset = ants_dataset+bees_dataset