随机对比度调整 随机裁剪pytorch代码

随机对比度调整：

import torch
import torchvision.transforms.functional as F

class RandomContrast(object):
    def __init__(self, lower=0.5, upper=1.5):
        self.lower = lower
        self.upper = upper

    def __call__(self, x):
        contrast_factor = torch.FloatTensor(1).uniform_(self.lower, self.upper)
        return F.adjust_contrast(x, contrast_factor)

随机裁剪：

import random
import torchvision.transforms.functional as F

class RandomCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, x):
        w, h = x.size
        th, tw = self.size
        if w == tw and h == th:
            return x
        i = random.randint(0, h - th)
        j = random.randint(0, w - tw)
        return F.crop(x, i, j, th, tw)

使用方式：

transform = transforms.Compose([
    RandomCrop(224),
    RandomContrast(),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

dataset = datasets.ImageFolder(root='path/to/data', transform=transform)