pytorch创建自己的数据集并传入dataloader（图像）

在之前的学习中已知使用的都是pytorch的官网里有的数据集，但是考虑到今后的学习中使用自己的数据集是难以避免的，所以这次写了个程序创建自己的数据集并传入dataloader。

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


class Mydataset(Dataset): #获取文件夹里的图片，并返回图片和label

    def __init__(self,root_dir,label_dir,transform=None):
        self.root_dir = root_dir
        self.label_dir = label_dir #
        self.path = os.path.join(self.root_dir,self.label_dir)
        self.imgs_path = os.listdir(self.path)
        self.transform = transform

    def __getitem__(self, idx):
        img_name = self.imgs_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        img = Image.open(img_item_path).convert('RGB') #转换成RGB形式的三通道图像
        if self.transform is not None:
            img = self.transform(img)
        label = self.label_dir
        return img,label

    def __len__(self):
        return len(self.imgs_path)

root_dir = r"test-train/train"
label_dir = r"奥特曼" #分别获取奥特曼和不是奥特曼文件夹里的图片构建数据集
label_notdir = r"不是奥特曼"
transform =  transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])
aoteDdataset = Mydataset(root_dir,label_dir,transform=transform)
not_aoteDataset = Mydataset(root_dir,label_notdir,transform=transform)
AoteDataset = aoteDdataset + not_aoteDataset #将两个数据集加起来变成一个完整的训练数据集
dataloader = torch.utils.data.DataLoader( AoteDataset, batch_size=4, shuffle=True ) #将训练集传入dataloader

for imgs,targets in dataloader: #打印label值
    print(targets[0])

结果：

说明：

以我的数据集为例：

分为训练集train和测试集test两个文件夹，train文件夹里分为奥特曼和不是奥特曼为label的两个文件夹，是奥特曼文件夹里放着一些奥特曼图片。

 

 那么我们想要获得train_dataloader的话我们要将奥特曼文件夹和不是奥特曼文件夹的图片加到一起，root_dir就是train文件夹路径，label_dir就是'奥特曼'（因为我们的label就是'奥特曼'和'不是奥特曼'两个，并且我们后面要将root_dir和label_dir加起来，所以label_dir仅仅是'奥特曼'即可）