torch Dataset类

Posted 2021-08-11 ʚVVcatɞ

实现Dataset类：

• __init__：初始化信息

• __getitem__：传入的图片路径和下标返回目录下的图片信息

• __len__：返回图片的张数

from torch.utils.data import Dataset
import os
import cv2 as cv

class MyData(Dataset):  # 继承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, index):
        img_name = self.img_path[index]  # 获取每一张图片的名称
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)  # 将存放图片的路径进行拼接
        img = cv.imread(img_item_path)   # 读取路径中的图片
        label = self.label_dir   # 图片的标签
        return img, label

    def __len__(self):
        return len(self.img_path)  # 获取有多少张图片


root_dir = "dataset/val"  # 根目录
ants_label_dir = "ants"   # 存放蚂蚁的目录
bees_label_dir = "bees"   # 存放蜜蜂的目录
ants_dataset = MyData(root_dir, ants_label_dir)  # 实例化 ants 
bees_label_dir = MyData(root_dir, bees_label_dir)  # 实例化 bees 

print("ants_dataset：", len(ants_dataset))
print("bees_label_dir：", len(bees_label_dir))

train_dataset = ants_dataset + bees_label_dir  # 将两个数据集存放到一起
print("train_dataset：", len(train_dataset))

img1, label1 = train_dataset[69]  # 查看第70张图片和标签信息

img2, label2 = train_dataset[70]  # 查看第71张图片和标签信息

#  图像显示
cv.imshow("image1", img1)
cv.imshow("image2", img2)
cv.waitKey(0)
cv.destroyAllWindows()