如何进行自定义 torchvision 变换？

Question

我有一个以 20% 的几率更改图像像素的函数，但不确定如何使其在 Transforms.Compose([]) 中工作。请帮忙！

def random_t(img):
    im = Image.open(img)
    pixelMap = im.load()
    pixelMap_list = []

    for i in range(im.size[0]):
        for j in range(im.size[1]):
            randNum = random.uniform(0, 1)
            if randNum < 0.2:  # 20% chance of pixel change
                pixelMap[i, j] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
            pixelMap_list.append(pixelMap[i, j])
        else:
            pixelMap[i, j] = pixelMap[i, j]

    return im

我认为它应该有这样的格式..这是来自 pytorch 库。

class custom_augmentation(object):
    def __init__(self, p):
        self.p = p # it should be the probability of random pixel

    def __call__(self, img):
        return None # Not sure how to make random_t in here

    def __repr__(self):
        return "custom augmentation"

固定代码：

class custom_augmentation(object):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img):
        pixelMap = img.load()
        for i in range(img.size[0]):
            for j in range(img.size[1]):
                if torch.rand(1) < self.p:
                    pixelMap[i, j] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
                else:
                    pixelMap[i, j] = pixelMap[i, j]

        return img # Not sure how to make random_t in here

    def __repr__(self):
        return "custom augmentation"

Answer 1

您需要对其进行操作img然后将其返回。有关如何创建自定义变换的一个很好的示例，只需查看如何创建正常的 torchvision 变换，如下所示：

这是 github，torchvision.transforms 之类的transforms.Resize()、transforms.ToTensor()、transforms.RandomHorizontalFlip()都有它们的代码。查看这些变换，例如RandomHorizontalFlip()了解如何引入变换发生的概率等。

https://github.com/pytorch/vision/blob/master/torchvision/transforms/transforms.py