PyTorch 中自定义数据集的读取方法

显然我们在学习深度学习时，不能只局限于通过使用官方提供的MNSIT、CIFAR-10、CIFAR-100这样的数据集，很多时候我们还是需要根据自己遇到的实际问题自己去搜集数据，然后制作数据集（收集数据集的方法有很多，这里就不过多的展开了）。这里只介绍数据集的读取。

1. 自定义数据集的方法：
   首先创建一个Dataset类
   
   在代码中：
   def init() 一些初始化的过程写在这个函数下
   def len() 返回所有数据的数量,比如我们这里将数据划分好之后，这里仅仅返回的是被处理后的关系
   def getitem() 回数据和标签

2. 补充代码
   上述已经将框架打出来了，接下来就是将框架填充完整就行了，下面是完整的代码，代码的解释说明我也已经写在其中了

# -*- coding: utf-8 -*-
# @Author  : 胡子旋
# @Email   ：1017190168@qq.com

import torch
import os,glob
import visdom
import time
import torchvision
import random,csv
from torch.utils.data import Dataset,DataLoader
from torchvision import transforms
from PIL import Image

class pokemom(Dataset):
def __init__(self,root,resize,mode,):
super(pokemom,self).__init__()
# 保存参数
self.root=root
self.resize=resize
# 给每一个类做映射
self.name2label={}  # "squirtle":0 ,"pikachu":1……
for name in sorted(os.listdir(os.path.join(root))):
# 过滤掉文件夹
if not os.path.isdir(os.path.join(root,name)):
continue
# 保存在表中;将最长的映射作为最新的元素的label的值
self.name2label[name]=len(self.name2label.keys())
print(self.name2label)
# 加载文件
self.images,self.labels=self.load_csv('images.csv')
# 裁剪数据
if mode=='train':
self.images=self.images[:int(0.6*len(self.images))]   # 将数据集的60%设置为训练数据集合
self.labels=self.labels[:int(0.6*len(self.labels))]   # label的60%分配给训练数据集合
elif mode=='val':
self.images = self.images[int(0.6 * len(self.images)):int(0.8 * len(self.images))]  # 从60%-80%的地方
self.labels = self.labels[int(0.6 * len(self.labels)):int(0.8 * len(self.labels))]
else:
self.images = self.images[int(0.8 * len(self.images)):]   # 从80%的地方到最末尾
self.labels = self.labels[int(0.8 * len(self.labels)):]
# image+label 的路径
def load_csv(self,filename):
# 将所有的图片加载进来
# 如果不存在的话才进行创建
if not os.path.exists(os.path.join(self.root,filename)):
images=[]
for name in self.name2label.keys():
images+=glob.glob(os.path.join(self.root,name,'*.png'))
images+=glob.glob(os.path.join(self.root, name, '*.jpg'))
images += glob.glob(os.path.join(self.root, name, '*.jpeg'))
print(len(images),images)
# 1167 'pokeman\\bulbasaur\\00000000.png'
# 将文件以上述的格式保存在csv文件内
random.shuffle(images)
with open(os.path.join(self.root,filename),mode='w',newline='') as f:
writer=csv.writer(f)
for img in images:    #  'pokeman\\bulbasaur\\00000000.png'
name=img.split(os.sep)[-2]
label=self.name2label[name]
writer.writerow([img,label])
print("write into csv into :",filename)

# 如果存在的话就直接的跳到这个地方
images,labels=[],[]
with open(os.path.join(self.root, filename)) as f:
reader=csv.reader(f)
for row in reader:
# 接下来就会得到 'pokeman\\bulbasaur\\00000000.png' 0 的对象
img,label=row
# 将label转码为int类型
label=int(label)
images.append(img)
labels.append(label)
# 保证images和labels的长度是一致的
assert len(images)==len(labels)
return images,labels

# 返回数据的数量
def __len__(self):
return len(self.images)   # 返回的是被裁剪之后的关系

def denormalize(self, x_hat):

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
# print(mean.shape, std.shape)
x = x_hat * std + mean
return x
# 返回idx的数据和当前图片的label
def __getitem__(self,idx):
# idex-[0-总长度]
# retrun images,labels
# 将图片，label的路径取出来
# 得到的img是这样的一个类型：'pokeman\\bulbasaur\\00000000.png'
# 然而label得到的则是 0，1，2 这样的整形的格式
img,label=self.images[idx],self.labels[idx]
tf=transforms.Compose([
lambda x:Image.open(x).convert('RGB'),  # 将t图片的路径转换可以处理图片数据
# 进行数据加强
transforms.Resize((int(self.resize*1.25),int(self.resize*1.25))),
# 随机旋转
transforms.RandomRotation(15),   # 设置旋转的度数小一些，否则的话会增加网络的学习难度
# 中心裁剪
transforms.CenterCrop(self.resize),   # 此时：既旋转了又不至于导致图片变得比较的复杂
transforms.ToTensor(),
transforms.Normalize(mean=[0.485,0.456,0.406],
std=[0.229,0.224,0.225])

])
img=tf(img)
label=torch.tensor(label)
return img,label

def main():
# 验证工作
viz=visdom.Visdom()

db=pokemom('pokeman',64,'train')  # 这里可以改变大小 224->64,可以通过visdom进行查看
# 可视化样本
x,y=next(iter(db))
print('sample:',x.shape,y.shape,y)
viz.image(db.denormalize(x),win='sample_x',opts=dict(title='sample_x'))
# 加载batch_size的数据
loader=DataLoader(db,batch_size=32,shuffle=True,num_workers=8)
for x,y in loader:
viz.images(db.denormalize(x),nrow=8,win='batch',opts=dict(title='batch'))
viz.text(str(y.numpy()),win='label',opts=dict(title='batch-y'))
# 每一次加载后，休息10s
time.sleep(10)

if __name__ == '__main__':
main()