PyTorch构造数据集（深度学习计算机视觉）

Posted 2021-08-31 SnoopyBUG

PyTorch构造数据集（深度学习计算机视觉）

• 构造数据集
• • 储存图片路径与标签的文件
  • 导入所需的库
  • 自定义datasets类
  • 定义load_data函数
  • 调用函数即可构造

构造数据集

在之前的学习过程中，训练模型用的数据集是pytorch自带的MNIST等，对导入数据集不求甚解。但当上手手势识别项目的时候，我对数据的导入一头雾水，查找到的资料让人“从入门到放弃”。自闭之后，重新研究了构造pytorch数据集的步骤，弄懂了其中一种相对简单的方法。

储存图片路径与标签的文件

以txt文件为例，train.txt和test.txt文件中每行都只有一张图片的路径、对应的标签。test.txt文件某片段如下，每行最后的数字是图片的标签。

test/2/hand1_2_top_seg_4_cropped.png 2
test/2/hand1_2_top_seg_3_cropped_0_9961.png 2
test/2/hand1_2_top_seg_3_cropped_0_7789.png 2
test/2/hand2_2_bot_seg_1_cropped.png 2
test/3/hand2_3_right_seg_1_cropped.png 3
test/3/hand5_3_bot_seg_2_cropped.png 3

可以先将图片根据标签储存在不同文件夹，在此基础上，使用python脚本创建这类txt文件。
（csv等文件亦可，本文以txt为例）

导入所需的库

import torch
import torchvision
from PIL import Image
import sys

自定义datasets类

自定义Dataset的子类——datasets类。
要之后要用Dataloader所定义的datasets类的话，这个类得拥有三个必要的函数：init、getitem、len。
注意：这里image_path在txt文件中path前加了’hand_gesture_data/’，保证image_path是正确的相对路径。

class gesture_datasets(torch.utils.data.Dataset):
    def __init__(self, txt_path, transform=None):
        lines = open(txt_path, 'r')
        imgs = []
        for line in lines:
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0], int(words[1])))
        self.imgs = imgs
        self.transform = transform
    def __getitem__(self, index):
        image_path = 'hand_gesture_data/' + self.imgs[index][0]
        label = self.imgs[index][1]
        img = Image.open(image_path).convert('RGB')
        img = self.transform(img)
        return img, label
    def __len__(self):
        return len(self.imgs)

定义load_data函数

这里定义了一个函数用于load data，其中将datasets类实例化后，关键在于torch.utils.data.DataLoader的使用。
txt_path同样是相对路径。使用Dataloader时可以考虑多线程（暂不支持Windows）。

def load_data_gesture(batch_size, resize=None):
    transform = []
    if resize:
        transform.append(torchvision.transforms.Resize(size=resize))
    transform.append(torchvision.transforms.ToTensor())
    transform = torchvision.transforms.Compose(transform)
    gesture_train = gesture_datasets(txt_path='hand_gesture_data/train/train.txt', transform=transform)
    gesture_test = gesture_datasets(txt_path='hand_gesture_data/test/test.txt', transform=transform)
    num_workers = 0 if sys.platform.startswith('win') else 4
    train_iter = torch.utils.data.DataLoader(gesture_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(gesture_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return train_iter, test_iter

调用函数即可构造

train_iter, test_iter = load_data_gesture(batch_size=60, resize=(224, 224))