How to rotate bounding box for tf$keras$preprocessing$image$apply_affine_transform

Question

I'm trying to write some functions for bounding box data augmentation. For image rotation I'm using tf$keras$preprocessing$image$apply_affine_transform.

Original image looks like this:

for which I have 3 bounding boxes (in YOLO format):

test_bbox <- structure(list(label = 1:3, xmin = c(0.236, 0.624, 0.001), ymin = c(0.396, 
0.592, 0.461333333333333), xmax = c(0.58, 0.966, 0.146), ymax = c(0.710666666666667, 
0.868, 0.670666666666667)), row.names = c(NA, -3L), class = "data.frame")

> test_bbox
  label  xmin      ymin  xmax      ymax
1     1 0.236 0.3960000 0.580 0.7106667
2     2 0.624 0.5920000 0.966 0.8680000
3     3 0.001 0.4613333 0.146 0.6706667

Plotting functions:

create_plot_data <- function(xy_axis, sample_image){
  cbind(xy_axis,
        r = as.vector(t(sample_image[, , 1])) / max(sample_image[, , 1]),
        g = as.vector(t(sample_image[, , 2])) / max(sample_image[, , 2]),
        b = as.vector(t(sample_image[, , 3])) / max(sample_image[, , 3]))
}

plot_rgb_raster <- function(plot_data){
  ggplot(plot_data, aes(x, y, fill = rgb(r, g, b))) +
    guides(fill = FALSE) + scale_fill_identity() + theme_void() +
    geom_raster(hjust = 0, vjust = 0)
}

correct_boxes <- function(boxes, image_h, image_w) {
  boxes %>% map(~ {
    current_boxes <- .x
    current_boxes %>%
      mutate(
        xmin = as.integer(xmin * image_w),
        ymin = as.integer(ymin * image_h),
        xmax = as.integer(xmax * image_w),
        ymax = as.integer(ymax * image_h)
      )
  })
}

apply_affine_transform_to_image <- function(image, theta = 0, tx = 0, ty = 0, shear = 0, zx = 1, zy = 1) {
  tf$keras$preprocessing$image$apply_affine_transform(image, theta, ty, tx, shear, zx, zy)
}

plot_boxes_ggplot <- function(image_path, boxes, correct_hw, target_size, theta = 0) {
  sample_image <- image_load(image_path, target_size = target_size) %>%
    image_to_array() %>% apply_affine_transform_to_image(theta = theta)
  h <- dim(sample_image)[1]
  w <- dim(sample_image)[2]
  boxes <- if (correct_hw) correct_boxes(list(boxes), image_h = h, image_w = w)[[1]] else boxes
  boxes <- boxes %>% mutate(x = 0, y = 0, r = 0, g = 0, b = 0)
  xy_axis <- expand.grid(1:w, h:1) %>% rename(x = Var1, y = Var2)
  plot_data <- create_plot_data(xy_axis, sample_image)
  p <- plot_rgb_raster(plot_data) +
    geom_rect(data = boxes, aes(xmin = xmin, ymin = h-ymin, xmax = xmax, ymax = h-ymax, colour = label),
              fill = NA, size = 1) +
    geom_label(data = boxes, aes(x = xmin, y = h-ymin, label = label, colour = label)) +
    theme(legend.position = "none")
  plot(p)
}

plot_boxes <- function(images_paths, boxes, correct_hw = TRUE, target_size = NULL) {
  walk2(images_paths, boxes, ~ plot_boxes_ggplot(.x, .y, correct_hw, target_size))
}

plot_boxes(images_paths = test_img, boxes = list(test_bbox), correct_hw = TRUE)

After that I wanted to rotate the image (I've noticed that apply_affine_transform is using positive degrees for clockwise rotation):

rotate_boxes <- function(boxes, height, width, theta = 0) {
  boxes %>%
    mutate_at(vars("xmin", "xmax"), ~ . * width) %>%
    mutate_at(vars("ymin", "ymax"), ~ . * height) %>%
    pmap_df(function(label, xmin, ymin, xmax, ymax) {
      x1 <- cos(-theta) * (xmin - width / 2) - sin(-theta) * (ymin - height / 2) + width / 2
      x2 <- cos(-theta) * (xmin - width / 2) - sin(-theta) * (ymax - height / 2) + width / 2
      x3 <- cos(-theta) * (xmax - width / 2) - sin(-theta) * (ymin - height / 2) + width / 2
      x4 <- cos(-theta) * (xmax - width / 2) - sin(-theta) * (ymax - height / 2) + width / 2
      y1 <- sin(-theta) * (xmin - width / 2) + cos(-theta) * (ymin - height / 2) + height / 2
      y2 <- sin(-theta) * (xmin - width / 2) + cos(-theta) * (ymax - height / 2) + height / 2
      y3 <- sin(-theta) * (xmax - width / 2) + cos(-theta) * (ymin - height / 2) + height / 2
      y4 <- sin(-theta) * (xmax - width / 2) + cos(-theta) * (ymax - height / 2) + height / 2
      tibble(label = label,
             xmin = min(c(x1, x2, x3, x4)) / width,
             ymin = min(c(y1, y2, y3, y4)) / height,
             xmax = max(c(x1, x2, x3, x4)) / width,
             ymax = max(c(y1, y2, y3, y4)) / height)
    })
}

plot_boxes(images_paths = test_img,
           boxes = list(rotate_boxes(test_bbox, 750, 1000, theta = 90)), correct_hw = TRUE, theta = 90)

I was trying different equations, but I can't figure out what's wrong.