Feat: add transform_to_rotated_box and box_xyxy_to_xyxyr

DESCRIPTION

@@ -83,6 +83,7 @@ Collate:
     'dataset-vggface2.R'
     'extension.R'
     'globals.R'
+    'item_transform_geometry.R'
     'models-alexnet.R'
     'models-convnext.R'
     'models-convnext_detection.R'

NAMESPACE

@@ -2,12 +2,15 @@
 
 S3method(draw_bounding_boxes,default)
 S3method(draw_bounding_boxes,image_with_bounding_box)
+S3method(draw_bounding_boxes,image_with_rotated_box)
 S3method(draw_bounding_boxes,torch_tensor)
 S3method(draw_segmentation_masks,default)
 S3method(draw_segmentation_masks,image_with_segmentation_mask)
 S3method(draw_segmentation_masks,torch_tensor)
 S3method(get_image_size,"magick-image")
 S3method(get_image_size,torch_tensor)
+S3method(item_transform_rotate,dataset)
+S3method(item_transform_rotate,image_with_bounding_box)
 S3method(prepare_sahi_split,"magick-image")
 S3method(prepare_sahi_split,dataset)
 S3method(prepare_sahi_split,numeric)
@@ -83,6 +86,7 @@ export(batched_nms)
 export(box_area)
 export(box_convert)
 export(box_iou)
+export(box_xyxy_to_xyxyr)
 export(caltech101_dataset)
 export(caltech256_dataset)
 export(caltech_classes)
@@ -115,6 +119,7 @@ export(imagenet_1k_classes)
 export(imagenet_21k_classes)
 export(imagenet_21k_df)
 export(imagenet_classes)
+export(item_transform_rotate)
 export(kmnist_dataset)
 export(lfw_pairs_dataset)
 export(lfw_people_dataset)

NEWS.md

@@ -12,6 +12,7 @@
 ## New features
 
 * Added `target_transform_resize` to manage bounding_box resizing in parallel to `transform_resize` for images (#337).
+* Added `item_transform_rotate()` for rotating both images and bounding boxes by an arbitrary angle (in degrees), returning boxes in `xyxyr` format. Includes a new `draw_bounding_boxes` S3 method for `image_with_rotated_box` that renders rotated boxes as polygons. (@DerrickUnleashed, #338)
 
 ## Bug fixes and improvements
 

R/item_transform_geometry.R

@@ -0,0 +1,70 @@
+#' Convert bounding boxes to rotated format
+#'
+#' Converts bounding boxes of a detection item from
+#' \eqn{(x_{min}, y_{min}, x_{max}, y_{max})} (xyxy) format to
+#' \eqn{(x_{min}, y_{min}, x_{max}, y_{max}, r)} (xyxyr) format, where
+#' \eqn{r} is the rotation angle in degrees (counter-clockwise). The image
+#' is left unchanged. For axis-aligned boxes, \eqn{r = 0}.
+#'
+#' @param x An object of class \code{image_with_bounding_box} or a dataset that
+#'   returns \code{image_with_bounding_box} items via \code{.getitem()}.
+#' @param angle (numeric): Rotation angle in degrees (counter-clockwise).
+#'   Default is \code{0}.
+#'
+#' @return An object of class \code{image_with_rotated_box} with the same
+#'   structure as the input, except that
+#'   \code{$y$boxes} is a tensor of shape \code{(N, 5)} in xyxyr format.
+#'   When applied to a dataset, returns the same dataset with its
+#'   \code{.getitem} method modified to return \code{image_with_rotated_box}
+#'   items.
+#'
+#' @examples
+#' \dontrun{
+#' url <- "https://upload.wikimedia.org/wikipedia/commons/9/9a/Aeroflot_Airbus_A330_Kustov.jpg"
+#'
+#' img <- base_loader(url) |>
+#'   transform_to_tensor() |>
+#'   transform_resize(c(300, 500))
+#'
+#' item <- list(
+#'   x = img,
+#'   y = list(
+#'     boxes = torch_tensor(matrix(c(40, 95, 475, 180), ncol = 4), dtype = torch_float32()),
+#'     labels = "airplane",
+#'     image_height = 300L,
+#'     image_width = 500L
+#'   )
+#' )
+#' class(item) <- c("image_with_bounding_box", "list")
+#'
+#' item_rot <- item_transform_rotate(item, angle = 355)
+#' item_rot$y$boxes  # (N, 5) tensor in xyxyr format
+#'
+#' before <- draw_bounding_boxes(item, colors = "blue", width = 4)
+#' after <- draw_bounding_boxes(item_rot, colors = "red", width = 4)
+#' tensor_image_browse(after)
+#' }
+#'
+#' @family item_unitary_transforms
+#'
+#' @export
+item_transform_rotate <- function(x, angle = 0) {
+  UseMethod("item_transform_rotate", x)
+}
+
+#' @export
+item_transform_rotate.image_with_bounding_box <- function(x, angle = 0) {
+  x$y$boxes <- box_xyxy_to_xyxyr(x$y$boxes, angle = angle)
+  class(x) <- c("image_with_rotated_box", setdiff(class(x), "image_with_bounding_box"))
+  x
+}
+
+#' @export
+item_transform_rotate.dataset <- function(x, angle = 0) {
+  original_getitem <- x$.getitem
+  x$.getitem <- function(index) {
+    item <- original_getitem(index)
+    item_transform_rotate(item, angle = angle)
+  }
+  x
+}

R/ops-box_convert.R

@@ -76,3 +76,73 @@ box_xyxy_to_xywh <- function(boxes) {
   return(boxes)
 }
 
+#' box_xyxy_to_xyxyr
+#'
+#' Converts bounding boxes from \eqn{(x_{min}, y_{min}, x_{max}, y_{max})} format to
+#'   \eqn{(x_{min}, y_{min}, x_{max}, y_{max}, r)} format, where \eqn{r} is the rotation
+#'   angle in degrees (anti-clockwise). For axis-aligned boxes, \eqn{r = 0}.
+#'
+#' @param boxes (Tensor\[N, 4\]): boxes in \eqn{(x_{min}, y_{min}, x_{max}, y_{max})} format
+#'   which will be converted.
+#' @param angle (numeric, optional): Rotation angle in degrees (anti-clockwise).
+#'   A single numeric value applied to all boxes, or a tensor of shape \code{(N,)}
+#'   with one angle per box. Default is \code{0}.
+#'
+#' @return boxes (Tensor\[N, 5\]): boxes in \eqn{(x_{min}, y_{min}, x_{max}, y_{max}, r)} format,
+#'   where \eqn{r} is the provided rotation angle in degrees. The bounding box
+#'   coordinates are computed by rotating the original axis-aligned box around
+#'   its center by \eqn{r} degrees anti-clockwise, then taking the axis-aligned
+#'   bounding box of the rotated corners.
+#'
+#' @export
+box_xyxy_to_xyxyr <- function(boxes, angle = 0) {
+
+  n <- boxes$size(1)
+
+  if (n == 0) {
+    angle_t <- torch_zeros(0, 1, dtype = boxes$dtype)
+    return(torch_cat(list(boxes, angle_t), dim = -1))
+  }
+
+  c(x1, y1, x2, y2) %<-% boxes$unbind(-1)
+  cx <- ((x1 + x2) / 2)$reshape(c(-1, 1))
+  cy <- ((y1 + y2) / 2)$reshape(c(-1, 1))
+  hw <- ((x2 - x1) / 2)$reshape(c(-1, 1))
+  hh <- ((y2 - y1) / 2)$reshape(c(-1, 1))
+
+  if (inherits(angle, "torch_tensor")) {
+    angle_deg <- angle$to(dtype = boxes$dtype)$reshape(c(-1, 1))
+  } else {
+    angle_deg <- torch_tensor(angle, dtype = boxes$dtype)$reshape(c(-1, 1))
+  }
+
+  if (angle_deg$size(1) == 1 && n > 1) {
+    angle_deg <- angle_deg$expand(c(n, 1))
+  }
+
+  angle_rad <- angle_deg * pi / 180
+  ct <- torch_cos(angle_rad)
+  st <- torch_sin(angle_rad)
+
+  corners_x <- torch_cat(list(
+    cx - hw * ct + hh * st,
+    cx + hw * ct + hh * st,
+    cx + hw * ct - hh * st,
+    cx - hw * ct - hh * st
+  ), dim = -1)
+
+  corners_y <- torch_cat(list(
+    cy - hw * st - hh * ct,
+    cy + hw * st - hh * ct,
+    cy + hw * st + hh * ct,
+    cy - hw * st + hh * ct
+  ), dim = -1)
+
+  xmin <- torch_min(corners_x, dim = -1)[[1]]$reshape(c(-1, 1))
+  xmax <- torch_max(corners_x, dim = -1)[[1]]$reshape(c(-1, 1))
+  ymin <- torch_min(corners_y, dim = -1)[[1]]$reshape(c(-1, 1))
+  ymax <- torch_max(corners_y, dim = -1)[[1]]$reshape(c(-1, 1))
+
+  torch_cat(list(xmin, ymin, xmax, ymax, angle_deg), dim = -1L)
+}
+

R/target_transform_detection.R

@@ -199,3 +199,5 @@ target_transform_sahi_crop <- function(y, sahi_split, min_area_ratio = 0.1) {
   results
 }
 
+
+

R/vision_utils.R

@@ -219,7 +219,98 @@ draw_bounding_boxes.image_with_bounding_box <- function(x, ...) {
   )
 }
 
+#' @rdname draw_bounding_boxes
+#' @export
+draw_bounding_boxes.image_with_rotated_box <- function(x, ...) {
+  rlang::check_installed("magick")
+  boxes <- x$y$boxes
+
+  img_to_draw <- if (x$x$dtype == torch::torch_uint8()) {
+    x$x$div(255)$permute(c(2, 3, 1))$to(device = "cpu") %>% as.array()
+  } else if (x$x$dtype == torch::torch_float()) {
+    x$x$permute(c(2, 3, 1))$to(device = "cpu") %>% as.array()
+  } else type_error("`x$x` should be torch_uint8 or torch_float")
+
+  num_boxes <- boxes$shape[1]
+  if (num_boxes == 0) {
+    cli_warn("boxes doesn't contain any box. No box was drawn")
+    return(x$x)
+  }
 
+  colors <- list(...)$colors
+  if (is.null(colors)) {
+    colors <- grDevices::hcl.colors(n = num_boxes)
+  }
+
+  labels <- x$y$labels
+  width <- list(...)$width
+  if (is.null(width)) width <- 1
+
+  font <- list(...)$font
+  if (is.null(font)) font <- c("serif", "plain")
+
+  font_size <- list(...)$font_size
+  if (is.null(font_size)) font_size <- 10
+
+  boxes_r <- as.matrix(boxes$to(device = "cpu"))
+
+  draw <- png::writePNG(img_to_draw) %>%
+    magick::image_read() %>%
+    magick::image_draw()
+
+  for (i in seq_len(num_boxes)) {
+    xmin <- boxes_r[i, 1]
+    ymin <- boxes_r[i, 2]
+    xmax <- boxes_r[i, 3]
+    ymax <- boxes_r[i, 4]
+    theta <- boxes_r[i, 5]
+
+    cx <- (xmin + xmax) / 2
+    cy <- (ymin + ymax) / 2
+    hw <- (xmax - xmin) / 2
+    hh <- (ymax - ymin) / 2
+
+    theta_rad <- theta * pi / 180
+    ct <- cos(theta_rad)
+    st <- sin(theta_rad)
+
+    corners_x <- c(
+      cx - hw * ct + hh * st,
+      cx + hw * ct + hh * st,
+      cx + hw * ct - hh * st,
+      cx - hw * ct - hh * st
+    )
+    corners_y <- c(
+      cy - hw * st - hh * ct,
+      cy + hw * st - hh * ct,
+      cy + hw * st + hh * ct,
+      cy - hw * st + hh * ct
+    )
+
+    graphics::polygon(corners_x, corners_y,
+                      border = colors[(i - 1) %% length(colors) + 1],
+                      lwd = width)
+
+    if (!is.null(labels)) {
+      lbl <- labels[(i - 1) %% length(labels) + 1]
+      graphics::text(corners_x[1] + 2 * width + font_size,
+                     corners_y[1] + 2 * width,
+                     labels = lbl,
+                     col = colors[(i - 1) %% length(colors) + 1],
+                     vfont = font,
+                     cex = font_size / 10)
+    }
+  }
+
+  grDevices::dev.off()
+
+  draw_tt <- draw %>%
+    magick::image_data(channels = "rgb") %>%
+    as.integer() %>%
+    torch::torch_tensor(dtype = torch::torch_uint8())
+
+  draw_tt$permute(c(3, 1, 2))
+}
 
 #' Convert COCO polygon to mask tensor (Robust Version)
 #'

_pkgdown.yml

@@ -33,30 +33,50 @@ navbar:
         href: articles/examples/keypoints.html
 
 reference:
-- title: Transforms
-  desc: Image transformation functions
-- subtitle: Unitary transformation
+- title: Image Transforms
+  desc: Transformation functions applied to Image (tensor images, magick images,...)
+- subtitle: Unitary image transformation
   contents:
     - has_concept("unitary_transforms")
-- subtitle: Random transformation
+- subtitle: Random image transformation
   contents:
     - has_concept("random_transforms")
-- subtitle: Combining / multiplying transformations
+- subtitle: Combining / multiplying image transformations
   contents:
     - has_concept("combining_transforms")
+
+- title: Target Transforms
+  desc: Transformation functions applied to targets (usualy the `$y` parameter)
 - subtitle: Target transformations for object detection
   desc: >
-    Functions to convert dataset native targets annotations into object-detection compatible
-    with draw_bounding_boxes() and object-detection models.
+    Functions to convert dataset native targets annotations into object-detection
+    targets compatible with draw_bounding_boxes() and object-detection models.
   contents:
     - has_concept("target_transforms_detection")
 - subtitle: Target transformations for segmentation
   desc: >
-    Functions to convert dataset native targets annotations into segmentation masks compatible
-    with draw_segmentation_masks() and segmentation models.
+    Functions to convert dataset native targets annotations into segmentation masks
+    targets compatible with draw_segmentation_masks() and segmentation models.
   contents:
     - has_concept("target_transforms_segmentation")
 
+- title: Item Transforms
+  desc: >
+    Transformation functions modifying both image (`$x` parameter ) and target
+    (`$y` parameter ) of a single dataset item. Result shall be compatible with
+    `draw_bounding_boxes()` for object-detection items and with `draw_segmentation_mask()`
+    for semantic-segmentation items.
+- subtitle: Unitary item transformations
+  desc: >
+    Functions to convert dataset **items** by an **unitary** transformation.
+  contents:
+    - has_concept("item_unitary_transforms")
+- subtitle: Random item transformations
+  desc: >
+    Functions to convert dataset **items** by a **random** transformation.
+  contents:
+    - has_concept("item_random_transforms")
+
 - title: Models
   desc: Computer Vision deep-learning Model architectures
 - subtitle: Classification models