20220117:技巧记录-在pytorch中应用albumentations

官网：Albumentations: fast and flexible image augmentationshttps://albumentations.ai/

在线可视化调参：不用自己慢慢写代码确定参数了

Streamlithttps://albumentations-demo.herokuapp.com/示例代码：

import os
import cv2
import torch
import numpy as np
from torch.utils.data import Dataset
import albumentations as A
from albumentations.pytorch import ToTensorV2



class AugDataset(Dataset):
    def __init__(self):
        super(AugDataset, self).__init__()


        transform_list = []

        # brightness
        aug_brightness = A.RandomBrightnessContrast(brightness_limit=0.18, contrast_limit=0.1, p=0.25)

        # noise
        iso_noise = A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.1, 0.5), p=0.5)
        gauss_noise = A.GaussNoise(var_limit=(5.0, 30.0), mean=0, p=0.5)
        aug_noise = A.OneOf([iso_noise, gauss_noise], p=0.5)
        transform_list.append(aug_noise)

        # blur
        med_blur = A.MedianBlur(blur_limit=(3, 5), p=0.25)
        moti_blur = A.MotionBlur(blur_limit=(5, 7), p=0.25)
        blur = A.Blur(blur_limit=(1, 5), p=0.25)
        gauss_blur = A.GaussianBlur(blur_limit=(1, 7), sigma_limit=0, p=0.25)
        aug_blur = A.OneOf([med_blur, moti_blur, blur, gauss_blur], p=0.5)
        transform_list.append(aug_blur)

        # compress
        #compress_a = A.ImageCompression(quality_lower=70, quality_upper=90, p=1.0)
        #scale_a = A.Downscale(scale_min=0.70, scale_max=0.90, interpolation=2, p=0.5)
        #aug_compress = A.OneOf([compress_a, scale_a], p=0.5)
        #transform_list.append(aug_compress)

        # choose one
        aug_choose = A.OneOf(transform_list, p=0.25)

        # all transform
        self.all_transform = A.Compose([
                aug_brightness,
                aug_choose,
                #A.HorizontalFlip(p=0.5),
                A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
                ToTensorV2(),
            ], p=1.0)


    def __getitem__(self, index):
        sample = cv2.imread(r'/home/lilai/workspace/Data/D2/process/20211209/buqueding/2307618_1_54.41_.jpg')

        sample = cv2.cvtColor(sample, cv2.COLOR_BGR2RGB)

        if self.all_transform is not None:
             sample = self.all_transform(image=sample)['image']

        return sample


    def __len__(self):
        return 100



if __name__ == "__main__":

    dataset = AugDataset()
    dataset.all_transform = A.Compose([t for t in dataset.all_transform if not isinstance(t, (A.Normalize, ToTensorV2))])

    while True:
        fig = np.zeros( (112*3, 112*3, 3), dtype=np.uint8)
        for idx in range(9):
            image = dataset[idx]

            row_idx = idx // 3
            col_idx = idx % 3
            fig[row_idx*112:(row_idx+1)*112, col_idx*112:(col_idx+1)*112,:] = image[:,:,::-1] # to bgr
        cv2.imshow('test', fig)
        cv2.waitKey(0)
        #cv2.imwrite("/home/mnt/lilai/Moredian/code/insight_pfc_torchcode_hhc/arcface_torch_dist_new/datasets0913.png", fig)