Commit 4e6faa75 authored by anguelos's avatar anguelos Committed by Francisco Massa
Browse files

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.
parent 3b450572
...@@ -205,6 +205,14 @@ class Tester(unittest.TestCase): ...@@ -205,6 +205,14 @@ class Tester(unittest.TestCase):
assert result.size(2) == width assert result.size(2) == width
assert np.allclose(img.numpy(), result.numpy()) 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): def test_pad(self):
height = random.randint(10, 32) * 2 height = random.randint(10, 32) * 2
width = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2
......
...@@ -368,14 +368,17 @@ class RandomCrop(object): ...@@ -368,14 +368,17 @@ class RandomCrop(object):
of the image. Default is 0, i.e no padding. If a sequence of length 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 4 is provided, it is used to pad left, top, right, bottom borders
respectively. 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): if isinstance(size, numbers.Number):
self.size = (int(size), int(size)) self.size = (int(size), int(size))
else: else:
self.size = size self.size = size
self.padding = padding self.padding = padding
self.pad_if_needed = pad_if_needed
@staticmethod @staticmethod
def get_params(img, output_size): def get_params(img, output_size):
...@@ -408,6 +411,13 @@ class RandomCrop(object): ...@@ -408,6 +411,13 @@ class RandomCrop(object):
if self.padding > 0: if self.padding > 0:
img = F.pad(img, self.padding) 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) i, j, h, w = self.get_params(img, self.size)
return F.crop(img, i, j, h, w) return F.crop(img, i, j, h, w)
......
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