@@ -368,14 +368,17 @@ class RandomCrop(object):
368
368
of the image. Default is 0, i.e no padding. If a sequence of length
369
369
4 is provided, it is used to pad left, top, right, bottom borders
370
370
respectively.
371
+ pad_if_needed (boolean): It will pad the image if smaller than the
372
+ desired size to avoid raising an exception.
371
373
"""
372
374
373
- def __init__ (self , size , padding = 0 ):
375
+ def __init__ (self , size , padding = 0 , pad_if_needed = False ):
374
376
if isinstance (size , numbers .Number ):
375
377
self .size = (int (size ), int (size ))
376
378
else :
377
379
self .size = size
378
380
self .padding = padding
381
+ self .pad_if_needed = pad_if_needed
379
382
380
383
@staticmethod
381
384
def get_params (img , output_size ):
@@ -408,6 +411,11 @@ def __call__(self, img):
408
411
if self .padding > 0 :
409
412
img = F .pad (img , self .padding )
410
413
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
+
411
419
i , j , h , w = self .get_params (img , self .size )
412
420
413
421
return F .crop (img , i , j , h , w )
0 commit comments