A minor change for transforms.RandomCrop (#462)

* Made transorms.RandomCrop tolerate images smaller than the given size.

* Extended the tescase for transforms.RandomCrop

* Made the testcase test for owidth=width+1

* Fixed the one pixel pading and the testcase

* Fixed minor lintint errors. flake8 passes.

Author: anguelos

test/test_transforms.py

@@ -205,6 +205,14 @@ def test_random_crop(self):
         assert result.size(2) == width
         assert np.allclose(img.numpy(), result.numpy())
 
+        result = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.RandomCrop((height + 1, width + 1), pad_if_needed=True),
+            transforms.ToTensor(),
+        ])(img)
+        assert result.size(1) == height + 1
+        assert result.size(2) == width + 1
+
     def test_pad(self):
         height = random.randint(10, 32) * 2
         width = random.randint(10, 32) * 2

torchvision/transforms/transforms.py

@@ -368,14 +368,17 @@ class RandomCrop(object):
             of the image. Default is 0, i.e no padding. If a sequence of length
             4 is provided, it is used to pad left, top, right, bottom borders
             respectively.
+        pad_if_needed (boolean): It will pad the image if smaller than the
+            desired size to avoid raising an exception.
     """
 
-    def __init__(self, size, padding=0):
+    def __init__(self, size, padding=0, pad_if_needed=False):
         if isinstance(size, numbers.Number):
             self.size = (int(size), int(size))
         else:
             self.size = size
         self.padding = padding
+        self.pad_if_needed = pad_if_needed
 
     @staticmethod
     def get_params(img, output_size):
@@ -408,6 +411,13 @@ def __call__(self, img):
         if self.padding > 0:
             img = F.pad(img, self.padding)
 
+        # pad the width if needed
+        if self.pad_if_needed and img.size[0] < self.size[1]:
+            img = F.pad(img, (int((1 + self.size[1] - img.size[0]) / 2), 0))
+        # pad the height if needed
+        if self.pad_if_needed and img.size[1] < self.size[0]:
+            img = F.pad(img, (0, int((1 + self.size[0] - img.size[1]) / 2)))
+
         i, j, h, w = self.get_params(img, self.size)
 
         return F.crop(img, i, j, h, w)