将自己的数据集加载到dataloader中

最新推荐文章于 2024-06-06 15:40:32 发布

zRezin

最新推荐文章于 2024-06-06 15:40:32 发布

阅读量470

点赞数 11

文章标签：深度学习计算机视觉数据处理

本文链接：https://blog.csdn.net/shuaikang9864/article/details/138122959

版权

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 3

2 文件名和标签存入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_infos

4 使用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()))