add Affine geometry transformation (#387)

* add Affine geometry transformation

* Add some docs.

tensorpack/dataflow/imgaug/geometry.py

@@ -8,7 +8,7 @@ import math
 import cv2
 import numpy as np
 
-__all__ = ['Shift', 'Rotation', 'RotationAndCropValid']
+__all__ = ['Shift', 'Rotation', 'RotationAndCropValid', 'Affine']
 
 
 class Shift(ImageAugmentor):
@@ -150,3 +150,101 @@ class RotationAndCropValid(ImageAugmentor):
             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 int(np.round(wr)), int(np.round(hr))
+
+
+class Affine(ImageAugmentor):
+    """
+    Random affine transform of the image w.r.t to the image center.
+    Transformations involve:
+        - Translation ("move" image on the x-/y-axis)
+        - Rotation
+        - Scaling ("zoom" in/out)
+        - Shear (move one side of the image, turning a square into a trapezoid)
+    """
+
+    def __init__(self, scale=None, translate_frac=None, rotate_max_deg=0.0, shear=0.0,
+                 interp=cv2.INTER_LINEAR, border=cv2.BORDER_REPLICATE, border_value=0):
+        """
+        Args:
+            scale (tuple of 2 floats): scaling factor interval, e.g (a, b), then scale is
+                randomly sampled from the range a <= scale <= b. Will keep
+                original scale by default.
+            translate_frac (tuple of 2 floats): tuple of max abs fraction for horizontal
+                and vertical translation. For example translate_frac=(a, b), then horizontal shift
+                is randomly sampled in the range 0 < dx < img_width * a and vertical shift is
+                randomly sampled in the range 0 < dy < img_height * b. Will
+                not translate by default.
+            shear (float): max abs shear value in degrees between 0 to 180
+            interp: cv2 interpolation method
+            border: cv2 border method
+            border_value: cv2 border value for border=cv2.BORDER_CONSTANT
+        """
+        if scale is not None:
+            assert isinstance(scale, tuple) and len(scale) == 2, \
+                "Argument scale should be a tuple of two floats, e.g (a, b)"
+
+        if translate_frac is not None:
+            assert isinstance(translate_frac, tuple) and len(translate_frac) == 2, \
+                "Argument translate_frac should be a tuple of two floats, e.g (a, b)"
+
+        assert shear >= 0.0, "Argument shear should be between 0.0 and 180.0"
+
+        super(Affine, self).__init__()
+        self._init(locals())
+
+    def _get_augment_params(self, img):
+        if self.scale is not None:
+            scale = self._rand_range(self.scale[0], self.scale[1])
+        else:
+            scale = 1.0
+
+        if self.translate_frac is not None:
+            max_dx = self.translate_frac[0] * img.shape[1]
+            max_dy = self.translate_frac[1] * img.shape[0]
+            dx = np.round(self._rand_range(-max_dx, max_dx))
+            dy = np.round(self._rand_range(-max_dy, max_dy))
+        else:
+            dx = 0
+            dy = 0
+
+        if self.shear > 0.0:
+            shear = self._rand_range(-self.shear, self.shear)
+            sin_shear = math.sin(math.radians(shear))
+            cos_shear = math.cos(math.radians(shear))
+        else:
+            sin_shear = 0.0
+            cos_shear = 1.0
+
+        center = (img.shape[1::-1] * np.array((0.5, 0.5))) - 0.5
+        deg = self._rand_range(-self.rotate_max_deg, self.rotate_max_deg)
+
+        transform_matrix = cv2.getRotationMatrix2D(tuple(center), deg, scale)
+
+        # Apply shear :
+        if self.shear > 0.0:
+            m00 = transform_matrix[0, 0]
+            m01 = transform_matrix[0, 1]
+            m10 = transform_matrix[1, 0]
+            m11 = transform_matrix[1, 1]
+            transform_matrix[0, 1] = m01 * cos_shear + m00 * sin_shear
+            transform_matrix[1, 1] = m11 * cos_shear + m10 * sin_shear
+            # Add correction term to keep the center unchanged
+            tx = center[0] * (1.0 - m00) - center[1] * transform_matrix[0, 1]
+            ty = -center[0] * m10 + center[1] * (1.0 - transform_matrix[1, 1])
+            transform_matrix[0, 2] = tx
+            transform_matrix[1, 2] = ty
+
+        # Apply shift :
+        transform_matrix[0, 2] += dx
+        transform_matrix[1, 2] += dy
+        return transform_matrix
+
+    def _augment(self, img, transform_matrix):
+        ret = cv2.warpAffine(img, transform_matrix, img.shape[1::-1],
+                             flags=self.interp, borderMode=self.border, borderValue=self.border_value)
+        if img.ndim == 3 and ret.ndim == 2:
+            ret = ret[:, :, np.newaxis]
+        return ret
+
+    def _augment_coords(self, coords, param):
+        raise NotImplementedError()