Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Apr 6, 2017
6 changes: 4 additions & 2 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions test/test_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 18 additions & 12 deletions torchvision/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import numpy as np
import numbers
import types
import collections


class Compose(object):
Expand Down Expand Up @@ -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):
Expand Down