使用albumentations数据增强库加载自己的数据集

使用albumentations数据增强库加载自己的数据集

1、Albumentations数据增强库简介

首先Albumentations库链接奉上：albumentations
截止目前已经有** 9.4K star**，足以看出该库使用的频率及认可。
该库中包含了大量增强算法，相比于pytorch自带的数据增强更有选择性与多样性。
本博客主要介绍如何使用albumentations库进行数据增强，并应用到自己的数据集中。

2、写一个加载数据集的类

该类的作用是加载自己的数据集(数据增强后的)

from torchvision.datasets.folder import *
from matplotlib import pyplot as plt
import cv2
import numpy as np
import torch.utils.data as Data
import albumentations as A
from albumentations.pytorch import ToTensorV2

class MyFolder(DatasetFolder):
    def __init__(self, root: str, transform=None, target_transform=None, loader=default_loader, is_valid_file=None):
        super(MyFolder, self).__init__(root, loader, IMG_EXTENSIONS if is_valid_file is None else None,
                                          transform=transform,
                                          target_transform=target_transform,
                                          is_valid_file=is_valid_file)
        self.imgs = self.samples

    def __getitem__(self, index: int):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """
        path, target = self.samples[index]
        # sample = self.loader(path)
        sample = cv2.imread(path)
        sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)
        if self.transform is not None:
            sample = self.transform(image=sample)["image"]
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

3、训练集与验证集的数据增强与加载

3.1 训练集的数据增强

图像增强：

    train_transform = A.Compose(
        [
        #仅仅显示部分，读者可去官网根据所需添加。
            A.SmallestMaxSize(max_size=160),
            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),
            A.Resize(height=128, width=128),
            A.RandomCrop(height=128, width=128),
            A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
            ToTensorV2(),
        ])

训练数据集加载：

data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
image_path = os.path.join(data_root, "I:/data", "数据集名称")  # flower data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
train_dataset = MyFolder(root=os.path.join(image_path, "train"), transform=train_transform)
train_loader = Data.DataLoader(dataset=train_dataset, num_workers=nw, batch_size=batch_s

3.2 验证集数据增强与加载

验证集数据增强：

    val_transform = A.Compose(
        [
            # A.SmallestMaxSize(max_size=160),
            # A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),
            A.Resize(height=128, width=128),
            # A.RandomCrop(height=128, width=128),
            # A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.5),
            # A.RandomBrightnessContrast(p=0.5),
            A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
            ToTensorV2(),
        ])

验证集数据集加载：

val_Dataset = MyFolder(root=os.path.join(image_path, "val"), transform=val_transform)
validate_loader = torch.utils.data.DataLoader(dataset=val_Dataset, num_workers=nw, batch_size=batch_size, shuffle=True)

4、总结

具体运行该图像增强数据集加载的完整代码，可参考图像分类github。