Skip to content

Commit 9c5c211

Browse files
committed
Made transorms.RandomCrop tolerate images smaller than the given size.
1 parent 3b45057 commit 9c5c211

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

torchvision/transforms/transforms.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -368,14 +368,17 @@ class RandomCrop(object):
368368
of the image. Default is 0, i.e no padding. If a sequence of length
369369
4 is provided, it is used to pad left, top, right, bottom borders
370370
respectively.
371+
pad_if_needed (boolean): It will pad the image if smaller than the
372+
desired size to avoid raising an exception.
371373
"""
372374

373-
def __init__(self, size, padding=0):
375+
def __init__(self, size, padding=0, pad_if_needed=False):
374376
if isinstance(size, numbers.Number):
375377
self.size = (int(size), int(size))
376378
else:
377379
self.size = size
378380
self.padding = padding
381+
self.pad_if_needed = pad_if_needed
379382

380383
@staticmethod
381384
def get_params(img, output_size):
@@ -408,6 +411,11 @@ def __call__(self, img):
408411
if self.padding > 0:
409412
img = F.pad(img, self.padding)
410413

414+
if self.pad_if_needed and img.size[0] < self.size[0]:
415+
img = F.pad(img, int((self.size[0] - img.size[0]) / 2), 0)
416+
if self.pad_if_needed and img.size[1] < self.size[1]:
417+
img = F.pad(img,(0, int((self.size[1] - img.size[1]) / 2)))
418+
411419
i, j, h, w = self.get_params(img, self.size)
412420

413421
return F.crop(img, i, j, h, w)

0 commit comments

Comments
 (0)