Commit f6d49d8b authored by Huan Yang's avatar Huan Yang Committed by Francisco Massa
Browse files

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
parent de5dcb93
......@@ -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
......
......@@ -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
......
......@@ -6,6 +6,7 @@ from PIL import Image, ImageOps
import numpy as np
import numbers
import types
import collections
class Compose(object):
......@@ -115,29 +116,34 @@ 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):
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):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment