add rotate and RandomRotation to transforms (#303)

add rotate and RandomRotation to transforms

Author: daavoo

test/test_transforms.py

@@ -664,6 +664,56 @@ def test_linear_transformation(self):
         cov = np.dot(xwhite, xwhite.T) / x.size(0)
         assert np.allclose(cov, np.identity(1), rtol=1e-3)
 
+    def test_rotate(self):
+        x = np.zeros((100, 100, 3), dtype=np.uint8)
+        x[40, 40] = [255, 255, 255]
+
+        with self.assertRaises(TypeError):
+            F.rotate(x, 10)
+
+        img = F.to_pil_image(x)
+
+        result = F.rotate(img, 45)
+        assert result.size == (100, 100)
+        r, c, ch = np.where(result)
+        assert all(x in r for x in [49, 50])
+        assert all(x in c for x in [36])
+        assert all(x in ch for x in [0, 1, 2])
+
+        result = F.rotate(img, 45, expand=True)
+        assert result.size == (142, 142)
+        r, c, ch = np.where(result)
+        assert all(x in r for x in [70, 71])
+        assert all(x in c for x in [57])
+        assert all(x in ch for x in [0, 1, 2])
+
+        result = F.rotate(img, 45, center=(40, 40))
+        assert result.size == (100, 100)
+        r, c, ch = np.where(result)
+        assert all(x in r for x in [40])
+        assert all(x in c for x in [40])
+        assert all(x in ch for x in [0, 1, 2])
+
+        result_a = F.rotate(img, 90)
+        result_b = F.rotate(img, -270)
+
+        assert np.all(np.array(result_a) == np.array(result_b))
+
+    def test_random_rotation(self):
+
+        with self.assertRaises(ValueError):
+            transforms.RandomRotation(-0.7)
+            transforms.RandomRotation([-0.7])
+            transforms.RandomRotation([-0.7, 0, 0.7])
+
+        t = transforms.RandomRotation(10)
+        angle = t.get_params(t.degrees)
+        assert angle > -10 and angle < 10
+
+        t = transforms.RandomRotation((-10, 10))
+        angle = t.get_params(t.degrees)
+        assert angle > -10 and angle < 10
+
 
 if __name__ == '__main__':
     unittest.main()

torchvision/transforms/functional.py

@@ -524,3 +524,28 @@ def adjust_gamma(img, gamma, gain=1):
 
     img = Image.fromarray(np_img, 'RGB').convert(input_mode)
     return img
+
+
+def rotate(img, angle, resample=False, expand=False, center=None):
+    """Rotate the image by angle and then (optionally) translate it by (n_columns, n_rows)
+
+
+    Args:
+        img (PIL Image): PIL Image to be rotated.
+        angle ({float, int}): In degrees degrees counter clockwise order.
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        expand (bool, optional): Optional expansion flag.
+            If true, expands the output image to make it large enough to hold the entire rotated image.
+            If false or omitted, make the output image the same size as the input image.
+            Note that the expand flag assumes rotation around the center and no translation.
+        center (2-tuple, optional): Optional center of rotation.
+            Origin is the upper left corner.
+            Default is the center of the image.
+    """
+    if not _is_pil_image(img):
+        raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
+
+    return img.rotate(angle, resample, expand, center)

torchvision/transforms/transforms.py

@@ -17,7 +17,7 @@
 
 __all__ = ["Compose", "ToTensor", "ToPILImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad",
            "Lambda", "RandomCrop", "RandomHorizontalFlip", "RandomVerticalFlip", "RandomResizedCrop",
-           "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation", "ColorJitter"]
+           "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation", "ColorJitter", "RandomRotation"]
 
 
 class Compose(object):
@@ -570,3 +570,62 @@ def __call__(self, img):
         transform = self.get_params(self.brightness, self.contrast,
                                     self.saturation, self.hue)
         return transform(img)
+
+
+class RandomRotation(object):
+    """Rotate the image by angle.
+
+    Args:
+        degrees (sequence or float or int): Range of degrees to select from.
+            If degrees is a number instead of sequence like (min, max), the range of degrees
+            will be (-degrees, +degrees).
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        expand (bool, optional): Optional expansion flag.
+            If true, expands the output to make it large enough to hold the entire rotated image.
+            If false or omitted, make the output image the same size as the input image.
+            Note that the expand flag assumes rotation around the center and no translation.
+        center (2-tuple, optional): Optional center of rotation.
+            Origin is the upper left corner.
+            Default is the center of the image.
+    """
+
+    def __init__(self, degrees, resample=False, expand=False, center=None):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            if len(degrees) != 2:
+                raise ValueError("If degrees is a sequence, it must be of len 2.")
+            self.degrees = degrees
+
+        self.resample = resample
+        self.expand = expand
+        self.center = center
+
+    @staticmethod
+    def get_params(degrees):
+        """Get parameters for ``rotate`` for a random rotation.
+
+        Returns:
+            sequence: params to be passed to ``rotate`` for random rotation.
+        """
+        angle = np.random.uniform(degrees[0], degrees[1])
+
+        return angle
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be rotated.
+
+        Returns:
+            PIL Image: Rotated image.
+        """
+
+        angle = self.get_params(self.degrees)
+
+        return F.rotate(img, angle, self.resample, self.expand, self.center)