Unified input for F.affine

Author: vfdev-5

test/test_functional_tensor.py

@@ -349,13 +349,20 @@ def test_resized_crop(self):
         )
 
     def test_affine(self):
-        # Let's do some tests on square image at first
+        # Tests on square image
         tensor, pil_img = self._create_data(26, 26)
+
+        scripted_affine = torch.jit.script(F.affine)
         # 1) identity map
         out_tensor = F.affine(tensor, angle=0, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
         self.assertTrue(
             tensor.equal(out_tensor), msg="{} vs {}".format(out_tensor[0, :5, :5], tensor[0, :5, :5])
         )
+        out_tensor = scripted_affine(tensor, angle=0, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
+        self.assertTrue(
+            tensor.equal(out_tensor), msg="{} vs {}".format(out_tensor[0, :5, :5], tensor[0, :5, :5])
+        )
+
         # 2) Test rotation
         test_configs = [
             (90, torch.rot90(tensor, k=1, dims=(-1, -2))),
@@ -367,29 +374,68 @@ def test_affine(self):
             (180, torch.rot90(tensor, k=2, dims=(-1, -2))),
         ]
         for a, true_tensor in test_configs:
-
-            out_tensor = F.affine(tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
-            if true_tensor is not None:
-                self.assertTrue(
-                    true_tensor.equal(out_tensor),
-                    msg="{}\n{} vs \n{}".format(a, out_tensor[0, :5, :5], true_tensor[0, :5, :5])
-                )
-            else:
-                true_tensor = out_tensor
-
-            out_pil_img = F.affine(pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
-            out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
-
-            num_diff_pixels = (true_tensor != out_pil_tensor).sum().item() / 3.0
-            ratio_diff_pixels = num_diff_pixels / true_tensor.shape[-1] / true_tensor.shape[-2]
-            # Tolerence : 6% of different pixels
-            self.assertLess(
-                ratio_diff_pixels,
-                0.06,
-                msg="{}\n{} vs \n{}".format(
-                    ratio_diff_pixels, true_tensor[0, :7, :7], out_pil_tensor[0, :7, :7]
+            for fn in [F.affine, scripted_affine]:
+                out_tensor = fn(tensor, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
+                if true_tensor is not None:
+                    self.assertTrue(
+                        true_tensor.equal(out_tensor),
+                        msg="{}\n{} vs \n{}".format(a, out_tensor[0, :5, :5], true_tensor[0, :5, :5])
+                    )
+                else:
+                    true_tensor = out_tensor
+
+                out_pil_img = F.affine(pil_img, angle=a, translate=[0, 0], scale=1.0, shear=[0.0, 0.0], resample=0)
+                out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
+
+                num_diff_pixels = (true_tensor != out_pil_tensor).sum().item() / 3.0
+                ratio_diff_pixels = num_diff_pixels / true_tensor.shape[-1] / true_tensor.shape[-2]
+                # Tolerance : less than 6% of different pixels
+                self.assertLess(
+                    ratio_diff_pixels,
+                    0.06,
+                    msg="{}\n{} vs \n{}".format(
+                        ratio_diff_pixels, true_tensor[0, :7, :7], out_pil_tensor[0, :7, :7]
+                    )
                 )
-            )
+        # 3) Test translation
+        test_configs = [
+            [10, 12], (12, 13)
+        ]
+        for t in test_configs:
+            for fn in [F.affine, scripted_affine]:
+                out_tensor = fn(tensor, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], resample=0)
+                out_pil_img = F.affine(pil_img, angle=0, translate=t, scale=1.0, shear=[0.0, 0.0], resample=0)
+                self.compareTensorToPIL(out_tensor, out_pil_img)
+
+        # 3) Test rotation + translation + scale + share
+        test_configs = [
+            (45, [5, 6], 1.0, [0.0, 0.0]),
+            (33, (5, -4), 1.0, [0.0, 0.0]),
+            (45, [5, 4], 1.2, [0.0, 0.0]),
+            (33, (4, 8), 2.0, [0.0, 0.0]),
+            (85, (10, -10), 0.7, [0.0, 0.0]),
+            (0, [0, 0], 1.0, [35.0, ]),
+            (25, [0, 0], 1.2, [0.0, 15.0]),
+            (45, [10, 0], 0.7, [2.0, 5.0]),
+            (45, [10, -10], 1.2, [4.0, 5.0]),
+        ]
+        for r in [0, ]:
+            for a, t, s, sh in test_configs:
+                for fn in [F.affine, scripted_affine]:
+                    out_tensor = fn(tensor, angle=a, translate=t, scale=s, shear=sh, resample=r)
+                    out_pil_img = F.affine(pil_img, angle=a, translate=t, scale=s, shear=sh, resample=r)
+                    out_pil_tensor = torch.from_numpy(np.array(out_pil_img).transpose((2, 0, 1)))
+
+                    num_diff_pixels = (out_tensor != out_pil_tensor).sum().item() / 3.0
+                    ratio_diff_pixels = num_diff_pixels / out_tensor.shape[-1] / out_tensor.shape[-2]
+                    # Tolerance : less than 5% of different pixels
+                    self.assertLess(
+                        ratio_diff_pixels,
+                        0.05,
+                        msg="{}: {}\n{} vs \n{}".format(
+                            (r, a, t, s, sh), ratio_diff_pixels, out_tensor[0, :7, :7], out_pil_tensor[0, :7, :7]
+                        )
+                    )
 
 
 if __name__ == '__main__':

test/test_transforms.py

@@ -1373,9 +1373,10 @@ def _test_transformation(a, t, s, sh):
             inv_true_matrix = np.linalg.inv(true_matrix)
             for y in range(true_result.shape[0]):
                 for x in range(true_result.shape[1]):
-                    res = np.dot(inv_true_matrix, [x, y, 1])
-                    _x = int(res[0] + 0.5)
-                    _y = int(res[1] + 0.5)
+                    # transform pixel's center instead of pixel's TL corner
+                    res = np.dot(inv_true_matrix, [x + 0.5, y + 0.5, 1])
+                    _x = int(res[0])
+                    _y = int(res[1])
                     if 0 <= _x < input_img.shape[1] and 0 <= _y < input_img.shape[0]:
                         true_result[y, x, :] = input_img[_y, _x, :]
 
@@ -1384,8 +1385,8 @@ def _test_transformation(a, t, s, sh):
             # Compute number of different pixels:
             np_result = np.array(result)
             n_diff_pixels = np.sum(np_result != true_result) / 3
-            # Accept 3 wrong pixels
-            self.assertLess(n_diff_pixels, 3,
+            # Accept 7 wrong pixels
+            self.assertLess(n_diff_pixels, 7,
                             "a={}, t={}, s={}, sh={}\n".format(a, t, s, sh) +
                             "n diff pixels={}\n".format(np.sum(np.array(result)[:, :, 0] != true_result[:, :, 0])))
 

torchvision/transforms/functional.py

@@ -791,7 +791,9 @@ def rotate(img, angle, resample=False, expand=False, center=None, fill=None):
     return img.rotate(angle, resample, expand, center, **opts)
 
 
-def _get_inverse_affine_matrix(center, angle, translate, scale, shear):
+def _get_inverse_affine_matrix(
+        center: List[int], angle: float, translate: List[float], scale: float, shear: List[float]
+) -> List[float]:
     # Helper method to compute inverse matrix for affine transformation
 
     # As it is explained in PIL.Image.rotate
@@ -818,14 +820,14 @@ def _get_inverse_affine_matrix(center, angle, translate, scale, shear):
     tx, ty = translate
 
     # RSS without scaling
-    a = cos(rot - sy) / cos(sy)
-    b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
-    c = sin(rot - sy) / cos(sy)
-    d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
+    a = math.cos(rot - sy) / math.cos(sy)
+    b = -math.cos(rot - sy) * math.tan(sx) / math.cos(sy) - math.sin(rot)
+    c = math.sin(rot - sy) / math.cos(sy)
+    d = -math.sin(rot - sy) * math.tan(sx) / math.cos(sy) + math.cos(rot)
 
     # Inverted rotation matrix with scale and shear
     # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
-    matrix = [d, -b, 0, -c, a, 0]
+    matrix = [d, -b, 0.0, -c, a, 0.0]
     matrix = [x / scale for x in matrix]
 
     # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
@@ -835,11 +837,12 @@ def _get_inverse_affine_matrix(center, angle, translate, scale, shear):
     # Apply center translation: C * RSS^-1 * C^-1 * T^-1
     matrix[2] += cx
     matrix[5] += cy
+
     return matrix
 
 
 def affine(
-        img: Tensor, angle: int, translate: List[int], scale: float, shear: List[float],
+        img: Tensor, angle: float, translate: List[int], scale: float, shear: List[float],
         resample: int = 0, fillcolor: Optional[int] = None
 ) -> Tensor:
     """Apply affine transformation on the image keeping image center invariant.
@@ -863,7 +866,10 @@ def affine(
     Returns:
         PIL Image or Tensor: Transformed image.
     """
-    if not isinstance(translate, Sequence):
+    if not isinstance(angle, (int, float)):
+        raise TypeError("Argument angle should be int or float")
+
+    if not isinstance(translate, (list, tuple)):
         raise TypeError("Argument translate should be a sequence")
 
     if len(translate) != 2:
@@ -872,30 +878,41 @@ def affine(
     if scale <= 0.0:
         raise ValueError("Argument scale should be positive")
 
-    if not isinstance(shear, (numbers.Number, Sequence)):
+    if not isinstance(shear, (numbers.Number, (list, tuple))):
         raise TypeError("Shear should be either a single value or a sequence of two values")
 
+    if isinstance(angle, int):
+        angle = float(angle)
+
+    if isinstance(translate, tuple):
+        translate = list(translate)
+
     if isinstance(shear, numbers.Number):
-        shear = [shear, 0]
+        shear = [shear, 0.0]
+
+    if isinstance(shear, tuple):
+        shear = list(shear)
+
+    if len(shear) == 1:
+        shear = [shear[0], shear[0]]
 
     if len(shear) != 2:
         raise ValueError("Shear should be a sequence containing two values. Got {}".format(shear))
 
+    img_size = _get_image_size(img)
     if not isinstance(img, torch.Tensor):
-        img_size = _get_image_size(img)
         # center = (img_size[0] * 0.5 + 0.5, img_size[1] * 0.5 + 0.5)
         # it is visually better to estimate the center without 0.5 offset
         # otherwise image rotated by 90 degrees is shifted vs output image of torch.rot90 or F_t.affine
-        center = (img_size[0] * 0.5, img_size[1] * 0.5)
+        center = [img_size[0] * 0.5, img_size[1] * 0.5]
         matrix = _get_inverse_affine_matrix(center, angle, translate, scale, shear)
 
         return F_pil.affine(img, matrix=matrix, resample=resample, fillcolor=fillcolor)
 
-    # compute affine matrix (not inversed)
-    # matrix = _get_inverse_affine_matrix(
-    #     (0, 0), -angle, [-t for t in translate], 1.0 / scale, [-s for s in shear]
-    # )
-    matrix = _get_inverse_affine_matrix((0, 0), angle, translate, scale, shear)
+    # we need to rescale translate by image size / 2 as its values can be between -1 and 1
+    translate = [2.0 * t / s for s, t in zip(img_size, translate)]
+
+    matrix = _get_inverse_affine_matrix([0, 0], angle, translate, scale, shear)
     return F_t.affine(img, matrix=matrix, resample=resample, fillcolor=fillcolor)
 
 

torchvision/transforms/functional_tensor.py

@@ -579,7 +579,7 @@ def resize(img: Tensor, size: List[int], interpolation: int = 2) -> Tensor:
 
 
 def affine(
-        img: Tensor, matrix: List[int], resample: int = 0, fillcolor: Optional[int] = None
+        img: Tensor, matrix: List[float], resample: int = 0, fillcolor: Optional[int] = None
 ) -> Tensor:
     """Apply affine transformation on the Tensor image keeping image center invariant.