coordinate mapping for affine transform

tensorpack/dataflow/imgaug/geometry.py

@@ -69,7 +69,7 @@ class Rotation(TransformAugmentorBase):
         if self.step_deg:
             deg = deg // self.step_deg * self.step_deg
         """
-        The correct center is shape*0.5-0.5 This can be verified by:
+        The correct center is shape*0.5-0.5. This can be verified by:
 
         SHAPE = 7
         arr = np.random.rand(SHAPE, SHAPE)

tensorpack/dataflow/imgaug/transform.py

@@ -102,5 +102,41 @@ class WarpAffineTransform(ImageTransform):
         return ret
 
     def apply_coords(self, coords):
-        # TODO
-        raise NotImplementedError()
+        coords = np.concatenate((coords, np.ones((coords.shape[0], 1), dtype='f4')), axis=1)
+        coords = np.dot(coords, self.mat.T)
+        return coords
+
+
+if __name__ == '__main__':
+    shape = (100, 100)
+    center = (10, 70)
+    mat = cv2.getRotationMatrix2D(center, 20, 1)
+    trans = WarpAffineTransform(mat, (130, 130))
+
+    def draw_points(img, pts):
+        for p in pts:
+            try:
+                img[int(p[1]), int(p[0])] = 0
+            except IndexError:
+                pass
+
+    image = cv2.imread('cat.jpg')
+    image = cv2.resize(image, shape)
+    orig_image = image.copy()
+    coords = np.random.randint(100, size=(20, 2))
+
+    draw_points(orig_image, coords)
+    print(coords)
+
+    for k in range(1):
+        coords = trans.apply_coords(coords)
+        image = trans.apply_image(image)
+    print(coords)
+    draw_points(image, coords)
+
+    #viz = cv2.resize(viz, (1200, 600))
+    orig_image = cv2.resize(orig_image, (600, 600))
+    image = cv2.resize(image, (600, 600))
+    viz = np.concatenate((orig_image, image), axis=1)
+    cv2.imshow("mat", viz)
+    cv2.waitKey()