Merge pull request #7 from pytorch/pad

adding padding to RandomCrop, as well as transforms.Pad

README.md

@@ -177,10 +177,11 @@ Crops the given PIL.Image at the center to have a region of
 the given size. 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)
 
-### `RandomCrop(size)`
+### `RandomCrop(size, padding=0)`
 Crops the given PIL.Image at a random location to have a region of
 the given size. 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)
+If `padding` is non-zero, then the image is first zero-padded on each side with `padding` pixels.
 
 ### `RandomHorizontalFlip()`
 Randomly horizontally flips the given PIL.Image with a probability of 0.5
@@ -193,6 +194,12 @@ This is popularly used to train the Inception networks
 - size: size of the smaller edge
 - interpolation: Default: PIL.Image.BILINEAR
 
+
+### `Pad(padding, fill=0)`
+Pads the given image on each side with `padding` number of pixels, and the padding pixels are filled with
+pixel value `fill`.
+If a `5x5` image is padded with `padding=1` then it becomes `7x7`
+
 ## Transforms on torch.*Tensor
 
 ### `Normalize(mean, std)`

test/test_transforms.py

@@ -81,6 +81,29 @@ class Tester(unittest.TestCase):
         assert result.size(1) == oheight
         assert result.size(2) == owidth
 
+        padding = random.randint(1, 20)
+        result = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.RandomCrop((oheight, owidth), padding=padding),
+            transforms.ToTensor(),
+        ])(img)
+        assert result.size(1) == oheight
+        assert result.size(2) == owidth
+
+    def test_pad(self):
+        height = random.randint(10, 32) * 2
+        width = random.randint(10, 32) * 2
+        img = torch.ones(3, height, width)
+        padding = random.randint(1, 20)
+        result = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Pad(padding),
+            transforms.ToTensor(),
+        ])(img)
+        print(height, width, padding)
+        print(result.size(1), result.size(2))
+        assert result.size(1) == height + 2*padding
+        assert result.size(2) == width + 2*padding
         
 
 if __name__ == '__main__':

torchvision/transforms.py

@@ -2,7 +2,7 @@ from __future__ import division
 import torch
 import math
 import random
-from PIL import Image
+from PIL import Image, ImageOps
 import numpy as np
 import numbers
 
@@ -115,6 +115,18 @@ class CenterCrop(object):
         return img.crop((x1, y1, x1 + tw, y1 + th))
 
 
+class Pad(object):
+    """Pads the given PIL.Image on all sides with the given "pad" value"""
+    def __init__(self, padding, fill=0):
+        assert isinstance(padding, numbers.Number)
+        assert isinstance(fill, numbers.Number)
+        self.padding = padding
+        self.fill = fill
+
+    def __call__(self, img):
+        return ImageOps.expand(img, border=self.padding, fill=self.fill)
+
+
 class RandomCrop(object):
     """Crops the given PIL.Image at a random location to have a region of
     the given size. size can be a tuple (target_height, target_width)
@@ -129,7 +141,7 @@ class RandomCrop(object):
 
     def __call__(self, img):
         if self.padding > 0:
-            raise NotImplementedError()
+            img = ImageOps.expand(img, border=self.padding, fill=0)
 
         w, h = img.size
         th, tw = self.size