Adds optional fill colour to rotate (#1280)

* Adds optional fill colour to rotate

* bug fix

Author: Philip Meier

torchvision/transforms/functional.py

@@ -686,7 +686,7 @@ def adjust_gamma(img, gamma, gain=1):
     return img
 
 
-def rotate(img, angle, resample=False, expand=False, center=None):
+def rotate(img, angle, resample=False, expand=False, center=None, fill=0):
     """Rotate the image by angle.
 
 
@@ -703,6 +703,8 @@ def rotate(img, angle, resample=False, expand=False, center=None):
         center (2-tuple, optional): Optional center of rotation.
             Origin is the upper left corner.
             Default is the center of the image.
+        fill (3-tuple or int): RGB pixel fill value for area outside the rotated image.
+            If int, it is used for all channels respectively.
 
     .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters
 
@@ -711,7 +713,10 @@ def rotate(img, angle, resample=False, expand=False, center=None):
     if not _is_pil_image(img):
         raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
 
-    return img.rotate(angle, resample, expand, center)
+    if isinstance(fill, int):
+        fill = tuple([fill] * 3)
+
+    return img.rotate(angle, resample, expand, center, fillcolor=fill)
 
 
 def _get_inverse_affine_matrix(center, angle, translate, scale, shear):

torchvision/transforms/transforms.py

@@ -946,12 +946,14 @@ class RandomRotation(object):
         center (2-tuple, optional): Optional center of rotation.
             Origin is the upper left corner.
             Default is the center of the image.
+        fill (3-tuple or int): RGB pixel fill value for area outside the rotated image.
+            If int, it is used for all channels respectively.
 
     .. _filters: https://pillow.readthedocs.io/en/latest/handbook/concepts.html#filters
 
     """
 
-    def __init__(self, degrees, resample=False, expand=False, center=None):
+    def __init__(self, degrees, resample=False, expand=False, center=None, fill=0):
         if isinstance(degrees, numbers.Number):
             if degrees < 0:
                 raise ValueError("If degrees is a single number, it must be positive.")
@@ -964,6 +966,7 @@ def __init__(self, degrees, resample=False, expand=False, center=None):
         self.resample = resample
         self.expand = expand
         self.center = center
+        self.fill = fill
 
     @staticmethod
     def get_params(degrees):
@@ -987,7 +990,7 @@ def __call__(self, img):
 
         angle = self.get_params(self.degrees)
 
-        return F.rotate(img, angle, self.resample, self.expand, self.center)
+        return F.rotate(img, angle, self.resample, self.expand, self.center, self.fill)
 
     def __repr__(self):
         format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)