Add support for Transforms.Scale([w, h]) with specific width and height (#133)

* Fixed border missing on bottom and right side using make_grid

* Add support for Transforms.Scale([h, w]) with specific height and width

if self.size is a int then scale the image with the shorter side,
otherwise if self.size is a list then scale the image to self.size
directly

* Add assert of size and doc for README

Add assert of size and doc for README

* Fix linter problem

Fix linter problem

* Add test for Scale

Add test for Scale

* Add both tuple and list support for Scale.size

Add both tuple and list support for Scale.size

* Add order of Scale.size in document and test case for list type of Scale.size

Add order of Scale.size in document and test case for list type of
Scale.size

README.rst

@@ -307,9 +307,11 @@ Transforms on PIL.Image
 ``Scale(size, interpolation=Image.BILINEAR)``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Rescales the input PIL.Image to the given 'size'. 'size' will be the
-size of the smaller edge.
+Rescales the input PIL.Image to the given 'size'. 
 
+If 'size' is a 2-element tuple or list in the order of (width, height), it will be the exactly size to scale.
+
+If 'size' is a number, it will indicate the size of the smaller edge. 
 For example, if height > width, then image will be rescaled to (size \*
 height / width, size) - size: size of the smaller edge - interpolation:
 Default: PIL.Image.BILINEAR

test/test_transforms.py

@@ -68,6 +68,24 @@ class Tester(unittest.TestCase):
         elif width < height:
             assert result.size(1) >= result.size(2)
 
+        oheight = random.randint(5, 12) * 2
+        owidth = random.randint(5, 12) * 2
+        result = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Scale((owidth, oheight)),
+            transforms.ToTensor(),
+        ])(img)
+        assert result.size(1) == oheight
+        assert result.size(2) == owidth
+
+        result = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Scale([owidth, oheight]),
+            transforms.ToTensor(),
+        ])(img)
+        assert result.size(1) == oheight
+        assert result.size(2) == owidth
+
     def test_random_crop(self):
         height = random.randint(10, 32) * 2
         width = random.randint(10, 32) * 2

torchvision/transforms.py

@@ -6,6 +6,7 @@ from PIL import Image, ImageOps
 import numpy as np
 import numbers
 import types
+import collections
 
 
 class Compose(object):
@@ -115,18 +116,21 @@ class Normalize(object):
 
 class Scale(object):
     """Rescales the input PIL.Image to the given 'size'.
-    'size' will be the size of the smaller edge.
+    If 'size' is a 2-element tuple or list in the order of (width, height), it will be the exactly size to scale.
+    If 'size' is a number, it will indicate the size of the smaller edge.
     For example, if height > width, then image will be
     rescaled to (size * height / width, size)
-    size: size of the smaller edge
+    size: size of the exactly size or the smaller edge
     interpolation: Default: PIL.Image.BILINEAR
     """
 
     def __init__(self, size, interpolation=Image.BILINEAR):
+        assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
         self.size = size
         self.interpolation = interpolation
 
     def __call__(self, img):
+        if isinstance(self.size, int):
             w, h = img.size
             if (w <= h and w == self.size) or (h <= w and h == self.size):
                 return img
@@ -138,6 +142,8 @@ class Scale(object):
                 oh = self.size
                 ow = int(self.size * w / h)
                 return img.resize((ow, oh), self.interpolation)
+        else:
+            return img.resize(self.size, self.interpolation)
 
 
 class CenterCrop(object):