from torch.utils.data import Dataset
class YourDataset(Dataset): # 继承Dataset类
# 构造函数必须存在
def __init__(self, root_dir, ann_file, transform=None):
self.ann_file = ann_file
self.root_dir = root_dir
self.img_label = self.load_annotations() # img_label是一个字典
self.img = [os.path.join(self.root_dir, img) for img in list(self.img_label.keys())]
self.label = [label for label in list(self.img_label.values())]
self.transform = transform # 数据需要做的预处理操作
def __len__(self):
return len(self.img)
# 获取图像和标签交给模型，该函数必须存在
# 不要修改参数，每次调用时会传入随机的idx
# 一个batch的数据就是由__getitem__函数处理数据传入得到的
def __getitem__(self, idx):
image = Image.open(self.img[idx]).convert('RGB') # img保存了图像的路径
label = self.label[idx]
if self.transform:
image = self.transform(image) # 对数据进行预处理操作
label = torch.from_numpy(np.array(label)) # 转换label的数据类型，由list->numpy->tensor
return image, label
def load_annotations(self):
data_infos = {}
with open(self.ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos前言
使用开源算法进行小规模的训练，例如分类任务需要加载自己的数据集和类别，需要使用dataloader格式的数据集。现在通过以下的方式将自己的数据集制作为dataloader格式：
参考
参考文档：深度学习(17)--DataLoader自定义数据集制作_自定义dataloader-CSDN博客
首先感谢分享，但是文档中有较多不明白的地方，并且也没有引用，在此上做了相对的改进：
实现过程
1 从txt文件中读取图片文件名和对应label
import numpy as np
def load_annotation(ann_file): # 参数为文本文件的路径
# 创建一个字典结构用于保存数据，key作为图像的名字，value作为图像的标签
data_infos = {}
with open(ann_file) as f:
# strip()去除一些换行符等
# split(' ')是以空格为分隔符
# samples是一个list，格式为图像名字，图像标签
# eg:[['image11.jpg,'0'],['image22.jpg,'1'],['image33.jpg,'3']]
samples = [x.strip().split(' ') for x in f.readlines()]
for filename, gt_label in samples:
# filename是图像名字--'image11.jpg'，gt_label--'0'是标签，加载到字典data_infos中去
# value值设置为array(gt_label,dtype=int64)类型
data_infos[filename] = np.array(gt_label, dtype=np.int64)
# 得到的字典格式：{'image11.jpg':array(0,dtype=int64),'image22.jpg':array(1,dtype=int64)}
return data_infos文件格式如下：
c1.jpg 1
c2.jpg 1
c3.jpg 1
c4.jpg 1
c5.jpg 1
l1.jpg 2
l2.jpg 2
l3.jpg 2
l4.jpg 2
l5.jpg 2
l6.jpg 2
q1.jpg 3
q2.jpg 3
q3.jpg 3
q4.jpg 3
q5.jpg 3
q6.jpg 32 文件名和标签存入list
img_label = load_annotation('testload.txt')
image_name = list(img_label.keys()) # 取keys值
label = list(img_label.values()) # 取labels值
print(img_label.keys()) # 参看你的数据3 整合为数据类
from torch.utils.data import Dataset
class YourDataset(Dataset): # 继承Dataset类
# 构造函数必须存在
def __init__(self, root_dir, ann_file, transform=None):
self.ann_file = ann_file
self.root_dir = root_dir
self.img_label = self.load_annotations() # img_label是一个字典
self.img = [os.path.join(self.root_dir, img) for img in list(self.img_label.keys())]
self.label = [label for label in list(self.img_label.values())]
self.transform = transform # 数据需要做的预处理操作
def __len__(self):
return len(self.img)
# 获取图像和标签交给模型，该函数必须存在
# 不要修改参数，每次调用时会传入随机的idx
# 一个batch的数据就是由__getitem__函数处理数据传入得到的
def __getitem__(self, idx):
image = Image.open(self.img[idx]).convert('RGB') # img保存了图像的路径
label = self.label[idx]
if self.transform:
image = self.transform(image) # 对数据进行预处理操作
label = torch.from_numpy(np.array(label)) # 转换label的数据类型，由list->numpy->tensor
return image, label
def load_annotations(self):
data_infos = {}
with open(self.ann_file) as f:
samples = [x.strip().split(' ') for x in f.readlines()]
for filename, gt_label in samples:
data_infos[filename] = np.array(gt_label, dtype=np.int64)
return data_infos4 使用transform函数进行数据处理
# 创建一个字典结构的数据类型来进行图像预处理操作：key - value
import torchvision.transforms as transforms
data_transforms = {
# 对训练集的预处理
'train': transforms.Compose([
transforms.Resize([256, 256]), # 卷积神经网络处理的数据大小必须相同，通过Resize来设置
# 数据增强
transforms.RandomRotation(45), # 随机旋转，-45到45度之间随机选
#transforms.CenterCrop(64), # 从中心开始裁剪，将原本96x96大小的图片数据裁剪为64x64大小的图片数据，可以获取更多的参数
transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转 选择一个概率概率，50%的概率进行水平翻转
transforms.RandomVerticalFlip(p=0.5), # 随 机垂直翻转，50%的概率进行竖直翻转
#transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1), # 参数1为亮度，参数2为对比度，参数3为饱和度，参数4为色相
#transforms.RandomGrayscale(p=0.025), # 概率转换成灰度率，3通道就是R=G=B(三颜色通道转为单一颜色通道，很少进行此处理)
# 将数据转为Tensor类型
transforms.ToTensor(),
# 标准化
transforms.Normalize([0.5, 0.5, 0.5], [0.224, 0.224, 0.225]) # 设置均值，标准差，分别对应R、G、B三个颜色通道的三个均值和标准差值，(x-μ)/σ
]),
# 对验证集的预处理(不需要进行数据增强)
'valid': transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# 均值和标准差数值的设置和训练集的相同(验证集的数据对我们来说是未知的，不能利用其中的数据再计算出相关的均值和标准差)
]),
}5 进行实例化
import os
from torch.utils.data import DataLoader
# 训练集
train_dataset = YourDataset(root_dir='.../...', ann_file='xxx.txt', transform=data_transforms['train'])
# 测试集
#valid_dataset = YourDataset(root_dir=valid_dir, ann_file='./.../valid.txt', transform=data_transforms['valid'])
# 实例化DataLoader(使用封装好的DataLoader包)
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
#valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=True)6 调用使用你的数据
dataloaders = {'train': train_loader}
#dataloaders = {'train': train_loader, "valid": valid_loader}
for inputs, labels in dataloaders['train']:
print("处理训练集")
#for inputs, labels in dataloaders['valid']:
# print("处理验证集")
#或者使用枚举
#for i, (imgs, labels) in enumerate(dataloaders['train']):
#测试完成后你可以将以上6步写入你的程序，或封装成数据读取包7 检查和展示你的数据
# 检查训练集
from PIL import Image
import torch
import matplotlib.pyplot as plt
image1, label1 = next(iter(train_loader)) # iter表示train_loader进行迭代，next取一个batch的数据
sample = image1[0].squeeze() # 通过squeeze()压缩一个维度，有时候维度为1x3x64x64，去除这个1
# 此时的sample是3x64x64的结构，而需要图像展示则需要转换结构为64X64X3，同时需要转换为numpy数据结构
sample = sample.permute((1, 2, 0)).numpy()
# 标准化还原 x = (x-μ) / σ -> x = x*σ + μ (预处理中进行了标准化，需要还原）
sample *= [0.229, 0.224, 0.225]
sample += [0.485, 0.456, 0.406]
plt.imshow(sample)
plt.show()
print('Label is: {}'.format(label1[0].numpy()))
# 检查训练集
#image2, label2 = next(iter(valid_loader)) # iter表示train_loader进行迭代，next取一个batch的数据
#sample = image2[0].squeeze() # 通过squeeze()压缩一个维度，有时候维度为1x3x64x64，去除这个1
# 此时的sample是3x64x64的结构，而需要图像展示则需要转换结构为64X64X3，同时需要转换为numpy数据结构
#sample = sample.permute((1, 2, 0)).numpy()
# 标准化还原 x = (x-μ) / σ -> x = x*σ + μ (预处理中进行了标准化，需要还原）
#sample *= [0.229, 0.224, 0.225]
#sample += [0.485, 0.456, 0.406]
#plt.imshow(sample)
#plt.show()
#print('Label is: {}'.format(label2[0].numpy()))
