Pytorch data enhancement method (code)

Data augmentation methods are often used to expand the size of the data set, improve the generalization ability, increase the robustness, and improve the classification accuracy.

Taking the input image as an example, the following code img is a [3,H,W] tensor, which can be run directly~

• random rotation
• Adjust image brightness, contrast, saturation, etc.
• add noise
• random cropping

import torch
import pandas as pd
import torchvision.transforms.functional as TF
from torchvision import transforms
import torchvision.transforms.functional as F
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
from torchvision import transforms
import json
import random

#数据增强
# 随机水平翻转
if random.random() > 0.5:
    img = F.hflip(img)
# 随机垂直翻转
if random.random() > 0.5:
    img = F.vflip(img)
# 随机旋转一定角度
angle = random.uniform(-10, 10)
img = F.rotate(img, angle)
# 随机对图像进行亮度、对比度和饱和度等色彩扰动
brightness_factor = 0.2
contrast_factor = 0.2
saturation_factor = 0.2
hue_factor = 0.2
img = F.adjust_brightness(img, brightness_factor)
img = F.adjust_contrast(img, contrast_factor)
img = F.adjust_saturation(img, saturation_factor)
img = F.adjust_hue(img, hue_factor)
# 向输入tensor添加高斯噪声
mean = 0
std = 1
noise = torch.randn(img.size()) * std + mean
img = img + noise
# 向输入tensor添加椒盐噪声
ratio = 0.05
mask = torch.rand(img.size())
img[mask < ratio] *= 0.1
#随机裁剪
C,H,W = img.size()
h_ran = random.uniform(0.8, 1.0)
w_ran = random.uniform(0.8, 1.0)
transform_crop = transforms.Compose([
    transforms.RandomCrop((int(H*h_ran), int(W*w_ran))),  # 先随机裁剪为稍小的大小
    transforms.Resize((H, W)),  # 再调整为原始大小
])
img = transform_crop(img)