Using PyTorch functional APIs for image data augmentation

img_augment.py

import torchvision.transforms.functional as F
import numpy as np
from PIL import Image
import os
from matplotlib import pyplot as plt

# input images
d = '../data/homes/test_input/'
img0 = Image.open(os.path.join(d, '352.jpg'))
img1 = Image.open(os.path.join(d, '95.jpg'))

# randomly rotate two images to the same angle
plt.subplot(221)
plt.imshow(img0)

plt.subplot(222)
plt.imshow(img1)

ang = np.random.randint(360)
print(ang)
img0_ = F.rotate(img0, ang)
img1_ = F.rotate(img1, ang)
    
plt.subplot(223)
plt.imshow(img0_)

plt.subplot(224)
plt.imshow(img1_)

# randomly flip two images to the same direction
plt.subplot(221)
plt.imshow(img0)

plt.subplot(222)
plt.imshow(img1)

if np.random.randint(2) == 0:
    img0_ = F.hflip(img0)
    img1_ = F.hflip(img1)
    
    plt.subplot(223)
    plt.imshow(img0_)

    plt.subplot(224)
    plt.imshow(img1_)
    
plt.subplot(221)
plt.imshow(img0)

plt.subplot(222)
plt.imshow(img1)

if np.random.randint(2) == 0:
    img0_ = F.vflip(img0)
    img1_ = F.vflip(img1)
    
    plt.subplot(223)
    plt.imshow(img0_)

    plt.subplot(224)
    plt.imshow(img1_)