pytorch自定义Transform

自定义一个随机加噪变换

值得注意的是pytorch中图像的格式都是PIL 的Image格式，如果用其他格式的进行处理，最后转回Image就可以了。

from torchvision import transforms
from PIL import Image
from torchvision.transforms import functional as TF
import torch
import skimage
import random
import numpy as np

# 读取一张测试图片
path = "./test.jpg"
img = Image.open(path)

def add_noise(img):
    a = random.random()
    if a <= 0.2:
        img=np.array(img)
        result = skimage.util.random_noise(img, mode='gaussian', seed=None, clip=True)
        return Image.fromarray(np.uint8(result*255))
    elif a > 0.2 and a <= 0.4:
        img=np.array(img)
        #随机生成5000个椒盐
        rows,cols,dims=img.shape
        for i in range(5000):
            x=np.random.randint(0,rows)
            y=np.random.randint(0,cols)
            img[x,y,:]=255
        for i in range(5000):
            x=np.random.randint(0,rows)
            y=np.random.randint(0,cols)
            img[x,y,:]=0
        return Image.fromarray(np.uint8(img))
    elif a > 0.4 and a <= 0.5:
        img=np.array(img)
        result = skimage.util.random_noise(img, mode='poisson', seed=None, clip=True)
        return Image.fromarray(np.uint8(result*255))
    # elif a > 0.6 and a <= 0.8:
    #     img=np.array(img)
    #     result = skimage.util.random_noise(img, mode='speckle', seed=None, clip=True)
    #     return Image.fromarray(np.uint8(result*255))
    else:
        return img
        
transform = transforms.Compose([
    transforms.Lambda(add_noise),
    # other transforms
])

new_img = transform(img)
new_img.save('./result.jpg')

自定义一个旋转变换

from torchvision import transforms
from PIL import Image
from torchvision.transforms import functional as TF
import torch
import skimage
import random
import numpy as np

# 读取一张测试图片
path = "./test.jpg"
img = Image.open(path)

class MyRotationTransform:
    """Rotate by one of the given angles."""

    def __init__(self, angles):
        self.angles = angles

    def __call__(self, x):
        angle = random.choice(self.angles)
        return TF.rotate(x, angle)

rotation_transform = MyRotationTransform(angles=[-30, -15, 0, 15, 30])

new_img = rotation_transform(img)
new_img.save('./result.jpg')