How to make a custom torchvision transform?

user_7422309 0
Viewed 2376

I have a function that changes image pixels with 20% chance, but not sure how to make it work in transforms.Compose([]). Please help!

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

I think it should have a format like this.. this is from pytorch library.

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"

fixed code:

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"
pytorch torchvision python
1 Answers

You need to do your operations on img and then return it. For a good example of how to create custom transforms just check out how the normal torchvision transforms are created like over here:

This is the github where torchvision.transforms like transforms.Resize(), transforms.ToTensor(), transforms.RandomHorizontalFlip() have their code. Look at these transforms like RandomHorizontalFlip() to see how to introduce a probability that a transform will happen etc.

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

Sarthak Jain

user_143470220
Related