fix transform_rotation using the image shape w × w instead of h × w (…

…#114) * fix rotation using the image shape w x w instead of h x w * add news --------- Co-authored-by: Daniel Falbel <[email protected]>
mlverse · Sep 17, 2024 · df06f28 · df06f28
1 parent bc2eab3
commit df06f28
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,6 @@
 # torchvision (development version)
 
+- fix `transform_rotation` wrongly using w x w for image size (@114, cregouby)
 - `tensor_image_display` and `tensor_image_browse` now accept all tensor_image dtypes. (#115, @cregouby) 
 - fix `transform_affine` help to remove confusion with `transforme_random_affine` help (#116, @cregouby)
 - add message translation in french (#112, @cregouby)

diff --git a/R/transforms-tensor.R b/R/transforms-tensor.R
@@ -663,8 +663,8 @@ rotate_impl <- function(img, matrix, resample = 0, expand = FALSE, fill= NULL) {
 
   assert_grid_transform_inputs(img, matrix, resample, fill, interpolation_modes)
   theta <- torch::torch_tensor(matrix)$reshape(c(1, 2, 3))
-  w <- tail(img$shape, 2)[2]
-  h <- tail(img$shape, 1)[1]
+  w <- tail(img$shape, 2)[1]
+  h <- tail(img$shape, 2)[2]
 
   if (expand) {
     o_shape <- rotate_compute_output_size(theta, w, h)

diff --git a/tests/testthat/test-transforms.R b/tests/testthat/test-transforms.R
@@ -134,7 +134,7 @@ test_that("rotate", {
   img <- torch::torch_tensor(matrix(1:16))$view(c(1, 4, 4))
   output <- transform_rotate(img, 90)
 
-  expect_tensor_shape(img, c(1,4,4))
+  expect_tensor_shape(output, c(1,4,4))
   expect_equal_to_r(output[1,,1], c(4,3,2,1))
 
   output <- transform_rotate(img, 45, expand = TRUE)
@@ -143,6 +143,17 @@ test_that("rotate", {
 
 })
 
+test_that("rotate a rectangle image", {
+
+  img <- torch::torch_tensor(matrix(1:20))$view(c(1, 4, 5))
+  output <- transform_rotate(img, 90)
+
+  expect_tensor_shape(output, c(1,5,4))
+  expect_equal_to_r(output[1,,1], c(5,4,3,2,1))
+  expect_equal_to_r(output[1,,4], c(20,19,18,17,16))
+
+})
+
 test_that("random_affine", {
 
   x <- torch_eye(8)$view(c(1, 1, 8, 8))