Feat: add transform_to_rotated_box and box_xyxy_to_xyxyr

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_bbox_rotate,dataset)
+S3method(item_transform_bbox_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_bbox_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_bbox_rotate()` for rotating bounding boxes by an arbitrary angle, 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/ops-box_convert.R

@@ -76,3 +76,72 @@ 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 radians (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 radians (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. The bounding box coordinates
+#'   are computed by rotating the original axis-aligned box around its center
+#'   by \eqn{r} radians 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 <- angle$to(dtype = boxes$dtype)$reshape(c(-1, 1))
+  } else {
+    angle <- torch_tensor(angle, dtype = boxes$dtype)$reshape(c(-1, 1))
+  }
+
+  if (angle$size(1) == 1 && n > 1) {
+    angle <- angle$expand(c(n, 1))
+  }
+
+  ct <- torch_cos(angle)
+  st <- torch_sin(angle)
+
+  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), dim = -1L)
+}
+

R/target_transform_detection.R

@@ -199,3 +199,61 @@ target_transform_sahi_crop <- function(y, sahi_split, min_area_ratio = 0.1) {
   results
 }
 
+#' Transform an image with bounding boxes to an image with rotated boxes
+#'
+#' Converts an \code{image_with_bounding_box} item (or a detection dataset that
+#' returns such items) to an \code{image_with_rotated_box} item. The bounding
+#' boxes are converted 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 radians (anti-clockwise). 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 radians (anti-clockwise). A
+#'   single numeric value applied to all boxes. Default is \code{0}.
+#'
+#' @return An object of class \code{image_with_rotated_box} with the same
+#'   structure as the input \code{image_with_bounding_box}, except that
+#'   \code{$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{
+#' # Convert a single item with a rotation angle
+#' ds <- coco_detection_dataset(train = FALSE, year = "2017", download = TRUE)
+#' item <- ds[1]
+#' rotated_item <- item_transform_bbox_rotate(item, angle = 0.5)
+#' rotated_item$y$boxes  # (N, 5) tensor in xyxyr format
+#'
+#' # Wrap a dataset with a rotation angle
+#' ds_rotated <- item_transform_bbox_rotate(ds, angle = 0.5)
+#' rotated_item <- ds_rotated[1]
+#' }
+#'
+#' @family target_transforms_detection
+#'
+#' @export
+item_transform_bbox_rotate <- function(x, angle = 0) {
+  UseMethod("item_transform_bbox_rotate", x)
+}
+
+#' @export
+item_transform_bbox_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", class(x))
+  x
+}
+
+#' @export
+item_transform_bbox_rotate.dataset <- function(x, angle = 0) {
+  original_getitem <- x$.getitem
+  x$.getitem <- function(index) {
+    item <- original_getitem(index)
+    item_transform_bbox_rotate(item, angle = angle)
+  }
+  x
+}
+

R/vision_utils.R

@@ -219,7 +219,97 @@ 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
+
+    ct <- cos(theta)
+    st <- sin(theta)
+
+    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)
 #'