diff --git a/test/test_transforms.py b/test/test_transforms.py index b3d55ae6b57..e2232e2491b 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -205,6 +205,14 @@ def test_random_crop(self): assert result.size(2) == width 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): height = random.randint(10, 32) * 2 width = random.randint(10, 32) * 2 diff --git a/torchvision/transforms/transforms.py b/torchvision/transforms/transforms.py index 8385b04df10..6475316074f 100644 --- a/torchvision/transforms/transforms.py +++ b/torchvision/transforms/transforms.py @@ -368,14 +368,17 @@ class RandomCrop(object): 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 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): self.size = (int(size), int(size)) else: self.size = size self.padding = padding + self.pad_if_needed = pad_if_needed @staticmethod def get_params(img, output_size): @@ -408,6 +411,13 @@ def __call__(self, img): if self.padding > 0: 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) return F.crop(img, i, j, h, w)