added shifts and stepped rotations (#177) (fix #170)

* added shifts and stepped rotations

* made shift random, fixed rotation to be pixel exact

* docs and linting

tensorpack/dataflow/imgaug/geometry.py

@@ -8,7 +8,36 @@ import math
 import cv2
 import numpy as np
 
-__all__ = ['Rotation', 'RotationAndCropValid']
+__all__ = ['Shift', 'Rotation', 'RotationAndCropValid']
+
+
+class Shift(ImageAugmentor):
+    """ Random horizontal and vertical shifts """
+
+    def __init__(self, horiz_frac=0, vert_frac=0,
+                 border=cv2.BORDER_REPLICATE):
+        """
+        Args:
+            horiz_frac (float): max abs fraction for horizontal shift
+            vert_frac (float): max abs fraction for horizontal shift
+            border: cv2 border method
+        """
+        assert horiz_frac < 1.0 and vert_frac < 1.0
+        super(Shift, self).__init__()
+        self._init(locals())
+
+    def _get_augment_params(self, img):
+        max_dx = self.horiz_frac * img.shape[1]
+        max_dy = self.vert_frac * img.shape[0]
+        dx = np.round(self._rand_range(-max_dx, max_dx))
+        dy = np.round(self._rand_range(-max_dy, max_dy))
+        return np.float32(
+            [[1, 0, dx], [0, 1, dy]])
+
+    def _augment(self, img, shift_m):
+        ret = cv2.warpAffine(img, shift_m, img.shape[1::-1],
+                             borderMode=self.border)
+        return ret
 
 
 class Rotation(ImageAugmentor):
@@ -16,14 +45,18 @@ class Rotation(ImageAugmentor):
 
     def __init__(self, max_deg, center_range=(0, 1),
                  interp=cv2.INTER_LINEAR,
-                 border=cv2.BORDER_REPLICATE):
+                 border=cv2.BORDER_REPLICATE, step_deg=None):
         """
         Args:
-            max_deg (float): max abs value of the rotation degree (in angle).
+            max_deg (float): max abs value of the rotation angle (in degree).
             center_range (tuple): (min, max) range of the random rotation center.
             interp: cv2 interpolation method
             border: cv2 border method
+            step_deg (float): if not None, the stepping of the rotation
+                angle. The rotation angle will be a multiple of step_deg. This
+                option requires ``max_deg==180`` and step_deg has to be a divisor of 180)
         """
+        assert step_deg is None or (max_deg == 180 and max_deg % step_deg == 0)
         super(Rotation, self).__init__()
         self._init(locals())
 
@@ -31,7 +64,9 @@ class Rotation(ImageAugmentor):
         center = img.shape[1::-1] * self._rand_range(
             self.center_range[0], self.center_range[1], (2,))
         deg = self._rand_range(-self.max_deg, self.max_deg)
-        return cv2.getRotationMatrix2D(tuple(center), deg, 1)
+        if self.step_deg:
+            deg = deg // self.step_deg * self.step_deg
+        return cv2.getRotationMatrix2D(tuple(center - 0.5), deg, 1)
 
     def _augment(self, img, rot_m):
         ret = cv2.warpAffine(img, rot_m, img.shape[1::-1],
@@ -44,22 +79,24 @@ class RotationAndCropValid(ImageAugmentor):
         Note that this will produce images of different shapes.
     """
 
-    def __init__(self, max_deg, interp=cv2.INTER_LINEAR):
+    def __init__(self, max_deg, interp=cv2.INTER_LINEAR, step_deg=None):
         """
         Args:
-            max_deg (float): max abs value of the rotation degree (in angle).
-            interp: cv2 interpolation method
+            max_deg, interp, step_deg: same as :class:`Rotation`
         """
+        assert step_deg is None or (max_deg == 180 and max_deg % step_deg == 0)
         super(RotationAndCropValid, self).__init__()
         self._init(locals())
 
     def _get_augment_params(self, img):
         deg = self._rand_range(-self.max_deg, self.max_deg)
+        if self.step_deg:
+            deg = deg // self.step_deg * self.step_deg
         return deg
 
     def _augment(self, img, deg):
         center = (img.shape[1] * 0.5, img.shape[0] * 0.5)
-        rot_m = cv2.getRotationMatrix2D(center, deg, 1)
+        rot_m = cv2.getRotationMatrix2D((center[0] - 0.5, center[1] - 0.5), deg, 1)
         ret = cv2.warpAffine(img, rot_m, img.shape[1::-1],
                              flags=self.interp, borderMode=cv2.BORDER_CONSTANT)
         neww, newh = RotationAndCropValid.largest_rotated_rect(ret.shape[1], ret.shape[0], deg)
@@ -95,5 +132,4 @@ class RotationAndCropValid(ImageAugmentor):
             # fully constrained case: crop touches all 4 sides
             cos_2a = cos_a * cos_a - sin_a * sin_a
             wr, hr = (w * cos_a - h * sin_a) / cos_2a, (h * cos_a - w * sin_a) / cos_2a
-
-        return np.round(wr), np.round(hr)
+        return int(np.round(wr)), int(np.round(hr))