shared_random_crop.py

from PIL import Image, ImageOps
import numbers
import random

class SharedRandomCrop(object):
    """Crops the given PIL.Image at a random location to have a region of
    the given size. Applies the same crop for a certain number (num_sharers) of
    consecutive crops. size can be a tuple (target_height, target_width)
    or an integer, in which case the target will be of a square shape (size, size)
    """
    def __init__(self, size, num_sharers, padding=0):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size
        self.num_sharers = num_sharers
        self._counter = 0
        self.padding = padding

    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
        # Get new crop
        if self._counter == 0:
            self.x1 = random.randint(0, w - tw)
            self.y1 = random.randint(0, h - th)
        self._counter = (self._counter + 1) % self.num_sharers
        print("cropping at {0}-{1}".format(self.x1,self.y1))
        return img.crop((self.x1, self.y1, self.x1 + tw, self.y1 + th))

# Test
if __name__ == "__main__":
    img = Image.new("RGB", (100,100))
    t = SharedRandomCrop(50, 2)
    for i in range(0,20):
        t(img)