PyTorch框架学习七——自定义transforms方法

虽然前面的笔记介绍了很多PyTorch给出的transforms方法，也非常有用，但是也有可能在具体的问题中需要开发者自定义transforms方法，这次笔记就介绍如何自定义transforms方法。

ps：本次笔记中使用的原始图像出自上次笔记：https://blog.csdn.net/qq_40467656/article/details/107958492

一、自定义transforms注意要素

从数据读取机制DataLoader我们知道了transforms的内部工作原理，是在Compose类的__call__函数定义并实现的：

class Compose(object):
    """Composes several transforms together.

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
    """

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

    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

由此出发，可以看出自定义transforms需要注意两个要素：

1. 仅接收一个参数img，并返回一个参数img；
2. transforms之间要注意上下游的输入与输出的格式匹配。

二、自定义transforms步骤

首先，自定义的transforms的输入参数可能不只img一个，如概率p等等，但是原来的代码只允许接收一个参数返回一个参数，所以可以在原来的基础上改进：

class YourTransforms(object):
	def __init__(self, ...):  # ...是要传入的多个参数
		# 对多参数进行传入
		# 如 self.p = p 传入概率
		# ...
	def __call__(self, img):  # __call__函数还是只有一个参数传入
		# 该自定义transforms方法的具体实现过程
		# ...
		return img

步骤如下：

1. 自定义一个类YourTransforms，结构类似Compose类
2. __init__函数作为多参数传入的地方
3. __call__函数具体实现自定义的transforms方法

三、自定义transforms实例：椒盐噪声

椒盐噪声：又称为脉冲噪声，是一种随机出现的白点或黑点，白点被称为盐噪声，黑点被称为椒噪声，其与信噪比（SNR）息息相关。

此外，我们还想加入概率p这个参数，实现随机添加椒盐噪声。

仿照实现步骤，先写出其实现的大致框架：

class AddPepperNoise(object):
	def __init__(self, snr, p):  # snr, p 是要传入的多个参数
		self.snr = snr
		self.p = p 
	def __call__(self, img):  # __call__函数还是只有一个参数传入
		'''
		添加椒盐噪声的具体实现过程
		'''
		return img

完整实现代码：

class AddPepperNoise(object):
    """增加椒盐噪声
    Args:
        snr （float）: Signal Noise Rate
        p (float): 概率值，依概率执行该操作
    """

    def __init__(self, snr, p=0.9):
        assert isinstance(snr, float) and (isinstance(p, float))
        self.snr = snr
        self.p = p

    def __call__(self, img):
        """
        Args:
            img (PIL Image): PIL Image
        Returns:
            PIL Image: PIL image.
        """
        if random.uniform(0, 1) < self.p:
            img_ = np.array(img).copy()
            h, w, c = img_.shape
            signal_pct = self.snr
            noise_pct = (1 - self.snr)
            mask = np.random.choice((0, 1, 2), size=(h, w, 1), p=[signal_pct, noise_pct/2., noise_pct/2.])
            mask = np.repeat(mask, c, axis=2)
            img_[mask == 1] = 255   # 盐噪声
            img_[mask == 2] = 0     # 椒噪声
            return Image.fromarray(img_.astype('uint8')).convert('RGB')
        else:
            return img

添加椒盐噪声之后：