Proposal for PyTorch Transforms

transforms.py

import collections
import numbers
import random
import math
from PIL import Image, ImageOps

def _iterate_transforms(transforms, x):
    if isinstance(transforms, collections.Iterable):
        for i, transform in enumerate(transforms):
            x[i] = _iterate_transforms(transform, x[i])
    else:
        x = transforms(x)
    return x

# we can pass nested arrays inside Compose
# the first level will be applied to all inputs
# and nested levels are passed to nested transforms
class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, x):
        for transform in self.transforms:
            x = _iterate_transforms(transform, x) 
        return x

class RandomCropGenerator(object):
    def __call__(self, img):
        self.x1 = random.uniform(0, 1)
        self.y1 = random.uniform(0, 1)
        return img

class RandomCrop(object):
    def __init__(self, size, padding=0, gen=None):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        self.padding = padding
        self._gen = gen

    def __call__(self, img):
        if self.padding > 0:
            img = ImageOps.expand(img, border=self.padding, fill=0)
        w, h = img.size
        th, tw = self.size
        if w == tw and h == th:
            return img

        if self._gen is not None:
            x1 = math.floor(self._gen.x1 * (w - tw))
            y1 = math.floor(self._gen.y1 * (h - th))
        else:
            x1 = random.randint(0, w - tw)
            y1 = random.randint(0, h - th)

        return img.crop((x1, y1, x1 + tw, y1 + th))

if __name__ == '__main__':
    path = '/Users/fmassa/workspace/vgg_face/data/Alba_Rohrwacher/00000783.jpg'
    im = Image.open(path)
    gen = RandomCropGenerator()
    t = Compose([
                gen,
                [RandomCrop(128, gen=gen), RandomCrop(128, gen=gen)]
            ])

    out = t([im, im])
    im.show()
    out[0].show()
    out[1].show()
    from IPython import embed; embed()