Add support for Transforms.Scale([w, h]) with specific width and height (pytorch#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

Author: yanghuanflc

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 @@ def test_scale(self):
         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 @@
 import numpy as np
 import numbers
 import types
+import collections
 
 
 class Compose(object):
@@ -115,29 +116,34 @@ def __call__(self, tensor):
 
 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):
-        w, h = img.size
-        if (w <= h and w == self.size) or (h <= w and h == self.size):
-            return img
-        if w < h:
-            ow = self.size
-            oh = int(self.size * h / w)
-            return img.resize((ow, oh), self.interpolation)
+        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
+            if w < h:
+                ow = self.size
+                oh = int(self.size * h / w)
+                return img.resize((ow, oh), self.interpolation)
+            else:
+                oh = self.size
+                ow = int(self.size * w / h)
+                return img.resize((ow, oh), self.interpolation)
         else:
-            oh = self.size
-            ow = int(self.size * w / h)
-            return img.resize((ow, oh), self.interpolation)
+            return img.resize(self.size, self.interpolation)
 
 
 class CenterCrop(object):