python - 如何进行自定义 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 转换，就像这里一样:

这是 torchvision.transforms 的 github有他们的代码。查看像 RandomHorizo​​ntalFlip() 这样的变换，了解如何引入变换发生的概率等。

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