用随机种子对image和mask做同样的数据增强

pytoch中的torchvision.transform是对输入数据按照一定概率做随机的变换，在神经网络中训练时候自然要保证input和target的一致性。

from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import random
import numpy as np
import os
import torch

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

alpha = Image.open('./train_alpha/0_0.png')
smoke = Image.open('./train_smoke/0_0.png')
image = Image.open('./train/0_0.png')


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

p1 = random.randint(0, 1)
p2 = random.randint(0, 1)
im_aug = transforms.Compose([
    transforms.RandomHorizontalFlip(p1),
    transforms.RandomVerticalFlip(p2),
    transforms.RandomRotation(10, resample=False, expand=False, center=None),
    transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
    transforms.RandomCrop(113),
    transforms.Resize((227, 227))
])
print(im_aug)

setup_seed(0)
f = plt.imshow(im_aug(alpha))
plt.show()

setup_seed(0)
f = plt.imshow(im_aug(smoke))
plt.show()

setup_seed(0)
f = plt.imshow(im_aug(image))
plt.show()

通过setup_seed()函数可以对alpha,smoke,image做同样的随机变换