【深度学习目标检测】Pytorch学习（1）DataSet数据集的加载

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


# 自定义类继承Dataset类
# 数据集的加载主要分为两个部分：一个是图片，一个为当前图片的标签
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(root_dir, 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 = "DataSet/hymenoptera_data/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