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
【深度学习目标检测】Pytorch学习(1)DataSet数据集的加载
猜你喜欢
转载自blog.csdn.net/beyond951/article/details/126271989
今日推荐
周排行