Merge pull request #99 from mickaelseznec/master

Fix for #97: handle correctly flow rotations

Author: ClementPinard

README.md

@@ -145,11 +145,10 @@ rotate: x,y,theta ->  (x*cos(theta)-x*sin(theta), y*cos(theta), x*sin(theta))
 ```
 
 #### Rotation applied on img2
-We consider the angle `theta` small enough to linearize `cos(theta)` to 1 and `sin(theta)` to `theta` .
+Let us consider a rotation by the angle `theta` from the image center.
 
-x flow map ( `flow[:,:,0]` ) will get a shift proportional to distance from center horizontal axis `j-h/2`
-
-y flow map ( `flow[:,:,1]` ) will get a shift proportional to distance from center vertical axis `i-w/2`
+We must tranform each flow vector based on the coordinates where it lands. On each coordinate `(i, j)`, we have:
 ```
-\for_all i,j flow[i,j] += theta*(j-h/2), theta*(i-w/2)
+flow[i, j, 0] += (cos(theta) - 1) * (j  - w/2 + flow[i, j, 0]) +    sin(theta)    * (i - h/2 + flow[i, j, 1])
+flow[i, j, 1] +=   -sin(theta)    * (j  - w/2 + flow[i, j, 0]) + (cos(theta) - 1) * (i - h/2 + flow[i, j, 1])
 ```

flow_transforms.py

@@ -167,7 +167,7 @@ class RandomRotate(object):
     angle: max angle of the rotation
     interpolation order: Default: 2 (bilinear)
     reshape: Default: false. If set to true, image size will be set to keep every pixel in the image.
-    diff_angle: Default: 0. Must stay less than 10 degrees, or linear approximation of flowmap will be off.
+    diff_angle: Default: 0.
     """
 
     def __init__(self, angle, diff_angle=0, order=2, reshape=False):
@@ -182,14 +182,22 @@ def __call__(self, inputs,target):
         angle1 = applied_angle - diff/2
         angle2 = applied_angle + diff/2
         angle1_rad = angle1*np.pi/180
+        diff_rad = diff*np.pi/180
 
         h, w, _ = target.shape
 
-        def rotate_flow(i,j,k):
-            return -k*(j-w/2)*(diff*np.pi/180) + (1-k)*(i-h/2)*(diff*np.pi/180)
+        warped_coords = np.mgrid[:w, :h].T + target
+        warped_coords -= np.array([w / 2, h / 2])
 
-        rotate_flow_map = np.fromfunction(rotate_flow, target.shape)
-        target += rotate_flow_map
+        warped_coords_rot = np.zeros_like(target)
+
+        warped_coords_rot[..., 0] = \
+            (np.cos(diff_rad) - 1) * warped_coords[..., 0] + np.sin(diff_rad) * warped_coords[..., 1]
+
+        warped_coords_rot[..., 1] = \
+            -np.sin(diff_rad) * warped_coords[..., 0] + (np.cos(diff_rad) - 1) * warped_coords[..., 1]
+
+        target += warped_coords_rot
 
         inputs[0] = ndimage.interpolation.rotate(inputs[0], angle1, reshape=self.reshape, order=self.order)
         inputs[1] = ndimage.interpolation.rotate(inputs[1], angle2, reshape=self.reshape, order=self.order)