"""
    Dataset: 创建数据集
"""
import os.path
import time

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


# 继承 Dataset 类
class MyData(Dataset):
    # 重写构造函数：为整个class提供全局变量
    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)
        print(self.img_path)

    # 使该类可以通过 对象[index] 获取指定元素
    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"  # 也可写成 "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)

img, label = ants_dataset[1]
img.show()  # 打开图片
print('图片标签：%s' % label)  # 查看标签

# 拼接两个数据集
train_dataset = ants_dataset + bees_dataset

time.sleep(2)