Fast rotation for right angles (#8295)

Co-authored-by: Thien Tran <thien.tran@parallelchain.io>

Author: gau-nernst

test/test_transforms_v2.py

@@ -1782,6 +1782,17 @@ def test_transform_unknown_fill_error(self):
  with pytest.raises(TypeError, match="Got inappropriate fill arg"):
  transforms.RandomAffine(degrees=0, fill="fill")
 
+ @pytest.mark.parametrize("size", [(11, 17), (16, 16)])
+ @pytest.mark.parametrize("angle", [0, 90, 180, 270])
+ @pytest.mark.parametrize("expand", [False, True])
+ def test_functional_image_fast_path_correctness(self, size, angle, expand):
+ image = make_image(size, dtype=torch.uint8, device="cpu")
+
+ actual = F.rotate(image, angle=angle, expand=expand)
+ expected = F.to_image(F.rotate(F.to_pil_image(image), angle=angle, expand=expand))
+
+ torch.testing.assert_close(actual, expected)
+
 
 class TestContainerTransforms:
  class BuiltinTransform(transforms.Transform):

torchvision/transforms/v2/functional/_geometry.py

@@ -997,6 +997,21 @@ def rotate_image(
  center: Optional[List[float]] = None,
  fill: _FillTypeJIT = None,
 ) -> torch.Tensor:
+ angle = angle % 360 # shift angle to [0, 360) range
+
+ # fast path: transpose without affine transform
+ if center is None:
+ if angle == 0:
+ return image.clone()
+ if angle == 180:
+ return torch.rot90(image, k=2, dims=(-2, -1))
+
+ if expand or image.shape[-1] == image.shape[-2]:
+ if angle == 90:
+ return torch.rot90(image, k=1, dims=(-2, -1))
+ if angle == 270:
+ return torch.rot90(image, k=3, dims=(-2, -1))
+
  interpolation = _check_interpolation(interpolation)
 
  input_height, input_width = image.shape[-2:]