CNN训练前的准备:PyTorch处理自己的图像数据(Dataset和Dataloader)
阅读原文时间:2023年08月09日阅读:1

目录

1. 数据处理

数据链接:cnn-dogs-vs-cats

PyTorch给我们提供了很多已经封装好的数据集,但是我们经常得使用自己找到的数据集,因此,想要得到一个好的训练结果,合理的数据处理是必不可少的。我们以1400张猫狗图片来进行分析。

  1. 分析数据:

训练集包含500张狗的图片以及500张猫的图片,测试集包含200张狗的图片以及200张猫的图片。

  1. 数据预处理:得到一个包含所有图片文件名(包含路径)和标签(狗1猫0)的列表:

    def init_process(path, lens):
    data = []
    name = find_label(path)
    for i in range(lens[0], lens[1]):
    data.append([path % i, name])

    return data

现有数据的命名都是有序号的,训练集中数据编号为0-499,测试集中编号为1000-1200,因此我们可以根据这个规律来读取文件名,比如参数传入:

path1 = 'cnn_data/data/training_data/cats/cat.%d.jpg'
data1 = init_process(path1, [0, 500])

data1就是一个包含五百个文件名以及标签的列表。find_label来判断标签是dog还是cat:

def find_label(str):
    first, last = 0, 0
    for i in range(len(str) - 1, -1, -1):
        if str[i] == '%' and str[i - 1] == '.':
            last = i - 1
        if (str[i] == 'c' or str[i] == 'd') and str[i - 1] == '/':
            first = i
            break

    name = str[first:last]
    if name == 'dog':
        return 1
    else:
        return 0

dog返回1,cat返回0。

有了上面两个函数之后,我们经过四次操作,就可以得到四个列表:

path1 = 'cnn_data/data/training_data/cats/cat.%d.jpg'
data1 = init_process(path1, [0, 500])
path2 = 'cnn_data/data/training_data/dogs/dog.%d.jpg'
data2 = init_process(path2, [0, 500])
path3 = 'cnn_data/data/testing_data/cats/cat.%d.jpg'
data3 = init_process(path3, [1000, 1200])
path4 = 'cnn_data/data/testing_data/dogs/dog.%d.jpg'
data4 = init_process(path4, [1000, 1200])

随便输出一个列表的前五个:

[['cnn_data/data/testing_data/dogs/dog.1000.jpg', 1], ['cnn_data/data/testing_data/dogs/dog.1001.jpg', 1], ['cnn_data/data/testing_data/dogs/dog.1002.jpg', 1], ['cnn_data/data/testing_data/dogs/dog.1003.jpg', 1], ['cnn_data/data/testing_data/dogs/dog.1004.jpg', 1]]
  1. 利用PIL包的Image处理图片:

    def Myloader(path):
    return Image.open(path).convert('RGB')

  2. 重写pytorch的Dataset类:

    class MyDataset(Dataset):
    def init(self, data, transform, loader):
    self.data = data
    self.transform = transform
    self.loader = loader
    def getitem(self, item):
    img, label = self.data[item]
    img = self.loader(img)
    img = self.transform(img)
    return img, label

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

里面有2个比较重要的函数:

  • __getitem__是真正读取数据的地方,迭代器通过索引来读取数据集中数据,因此只需要这一个方法中加入读取数据的相关功能即可。在这个函数里面,我们对第二步处理得到的列表进行索引,接着利用第三步定义的Myloader来对每一个路径进行处理,最后利用pytorch的transforms对RGB数据进行处理,将其变成Tensor数据。
    transform为:

    transform = transforms.Compose([
    transforms.CenterCrop(224),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) # 归一化
    ])

对上面四个操作做一些解释:
1)、transforms.CenterCrop(224),从图像中心开始裁剪图像,224为裁剪大小
2)、transforms.Resize((224, 224)),重新定义图像大小
3)、 transforms.ToTensor(),很重要的一步,将图像数据转为Tensor
4)、transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),归一化

  • __len__中提供迭代器索引的范围。

因此我们只需要:

# shuffle
np.random.shuffle(data)
# train, val, test = 900 + 200 + 300
train_data, val_data, test_data = data[:900], data[900:1100], data[1100:]
train_data = MyDataset(train_data, transform=transform, loader=Myloader)
Dtr = DataLoader(dataset=train_data, batch_size=50, shuffle=True, num_workers=0)
val_data = MyDataset(val_data, transform=transform, loader=Myloader)
Val = DataLoader(dataset=val_data, batch_size=50, shuffle=True, num_workers=0)
test_data = MyDataset(test_data, transform=transform, loader=Myloader)
Dte = DataLoader(dataset=test_data, batch_size=50, shuffle=True, num_workers=0)

就可以得到处理好的Dataset,其中训练集、验证集以及测试集分别有900张图片、200张图片以及300张图片。

最后我们只要给定义好的神经网络模型喂数据就OK了!!!

2. 完整代码

# -*- coding:utf-8 -*-
"""
@Time: 2022/03/01 11:33
@Author: KI
@File: data_process.py
@Motto: Hungry And Humble
"""
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np


def Myloader(path):
    return Image.open(path).convert('RGB')


# get a list of paths and labels.
def init_process(path, lens):
    data = []
    name = find_label(path)
    for i in range(lens[0], lens[1]):
        data.append([path % i, name])

    return data


class MyDataset(Dataset):
    def __init__(self, data, transform, loader):
        self.data = data
        self.transform = transform
        self.loader = loader

    def __getitem__(self, item):
        img, label = self.data[item]
        img = self.loader(img)
        img = self.transform(img)
        return img, label

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


def find_label(str):
    """
    Find image tags based on file paths.

    :param str: file path
    :return: image label
    """
    first, last = 0, 0
    for i in range(len(str) - 1, -1, -1):
        if str[i] == '%' and str[i - 1] == '.':
            last = i - 1
        if (str[i] == 'c' or str[i] == 'd') and str[i - 1] == '/':
            first = i
            break

    name = str[first:last]
    if name == 'dog':
        return 1
    else:
        return 0


def load_data():
    print('data processing...')
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(p=0.3),
        transforms.RandomVerticalFlip(p=0.3),
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))  # normalization
    ])
    path1 = 'data/training_data/cats/cat.%d.jpg'
    data1 = init_process(path1, [0, 500])
    path2 = 'data/training_data/dogs/dog.%d.jpg'
    data2 = init_process(path2, [0, 500])
    path3 = 'data/testing_data/cats/cat.%d.jpg'
    data3 = init_process(path3, [1000, 1200])
    path4 = 'data/testing_data/dogs/dog.%d.jpg'
    data4 = init_process(path4, [1000, 1200])
    data = data1 + data2 + data3 + data4   # 1400
    # shuffle
    np.random.shuffle(data)
    # train, val, test = 900 + 200 + 300
    train_data, val_data, test_data = data[:900], data[900:1100], data[1100:]
    train_data = MyDataset(train_data, transform=transform, loader=Myloader)
    Dtr = DataLoader(dataset=train_data, batch_size=50, shuffle=True, num_workers=0)
    val_data = MyDataset(val_data, transform=transform, loader=Myloader)
    Val = DataLoader(dataset=val_data, batch_size=50, shuffle=True, num_workers=0)
    test_data = MyDataset(test_data, transform=transform, loader=Myloader)
    Dte = DataLoader(dataset=test_data, batch_size=50, shuffle=True, num_workers=0)

    return Dtr, Val, Dte

train_data以及test_data就是我们最终需要得到的数据。

对猫狗数据分类的具体实现请见:CNN简单实战:PyTorch搭建CNN对猫狗图片进行分类

如果需要更高的准确率,可以使用一些预训练的模型,详见:
PyTorch搭建预训练AlexNet、DenseNet、ResNet、VGG实现猫狗图片分类