Feat: item transform random resize crop

NAMESPACE

@@ -34,6 +34,11 @@ S3method(item_transform_random_horizontal_flip,default)
 S3method(item_transform_random_horizontal_flip,image_with_bounding_box)
 S3method(item_transform_random_horizontal_flip,image_with_rotated_box)
 S3method(item_transform_random_horizontal_flip,image_with_segmentation_mask)
+S3method(item_transform_random_resize_crop,dataset)
+S3method(item_transform_random_resize_crop,default)
+S3method(item_transform_random_resize_crop,image_with_bounding_box)
+S3method(item_transform_random_resize_crop,image_with_rotated_box)
+S3method(item_transform_random_resize_crop,image_with_segmentation_mask)
 S3method(item_transform_random_vertical_flip,dataset)
 S3method(item_transform_random_vertical_flip,default)
 S3method(item_transform_random_vertical_flip,image_with_bounding_box)
@@ -172,6 +177,7 @@ export(item_transform_center_crop)
 export(item_transform_crop)
 export(item_transform_hflip)
 export(item_transform_random_horizontal_flip)
+export(item_transform_random_resize_crop)
 export(item_transform_random_vertical_flip)
 export(item_transform_pad)
 export(item_transform_rotate)
@@ -384,6 +390,7 @@ importFrom(torch,nnf_silu)
 importFrom(torch,nnf_softmax)
 importFrom(torch,torch_abs)
 importFrom(torch,torch_arange)
+importFrom(torch,torch_atan2)
 importFrom(torch,torch_bool)
 importFrom(torch,torch_cat)
 importFrom(torch,torch_chunk)

NEWS.md

@@ -26,9 +26,14 @@
 * `item_transform_rotate()` now supports segmentation items and datasets, rotating the masks alongside the image (@srishtiii28, #379).
 * Added `item_transform_crop()` for cropping dataset items at a specified location and size, with support for detection and segmentation item types and datasets (@DerrickUnleashed, #371).
 * Added `item_transform_pad()` for padding dataset items on all sides, with support for detection, segmentation and rotated-box item types and datasets (@DerrickUnleashed, #373).
+* Added `item_transform_random_resize_crop()` for cropping dataset items to a random area and aspect ratio before resizing them to a given size, with support for detection, segmentation and rotated-box item types and datasets (@srishtiii28, #361).
 
 ## Bug fixes and improvements
 
+* `transform_random_resized_crop()` no longer loses the last row and column of the image when it
+  falls back to a central crop. The fallback returned a zero-based corner while the sampled crops
+  return a one-based one, so the crop started one pixel above and left of the image and the result
+  was padded with a black edge (#361).
 * `transform_crop()` now pads the result with zeros when the crop leaves the image, so that the
   output always has the requested size. It previously returned only the part of the image the crop
   covered, which could even be empty.

R/item-transforms-random-geometry.R

@@ -165,3 +165,153 @@ item_transform_random_vertical_flip.image_with_rotated_box <- function(x, p = 0.
   }
   x
 }
+
+#' Randomly crop a dataset item and resize it
+#'
+#' Crops the image inside a dataset item to a random area and aspect ratio, then
+#' resizes that crop to the given size. For detection items, bounding boxes are
+#' cropped along with the image and their coordinates are rescaled to the output
+#' size: boxes falling outside the crop are dropped and boxes straddling its
+#' border are clipped to it. For segmentation items, the masks are cropped and
+#' resized alongside the image with nearest-neighbour sampling, so that they keep
+#' their discrete values.
+#'
+#' The crop is drawn again for every item, so that a dataset wrapped with this
+#' transform yields a different crop on each access.
+#'
+#' The \code{area} field of a detection target, when present, is rescaled by the
+#' area ratio between the crop and the output. For rotated boxes, a crop whose
+#' aspect ratio differs from the output one maps the box to a parallelogram: the
+#' enclosing box stays exact and the angle is that of the transformed box axis,
+#' which is exact whenever both edges are scaled alike.
+#'
+#' @param x A dataset item, typically an \code{image_with_bounding_box} or
+#'   \code{image_with_segmentation_mask} object containing an image tensor
+#'   and associated target data.
+#' @inheritParams transform_random_resized_crop
+#'
+#' @return A dataset item of the same class with the image and target cropped
+#'   and resized.
+#'
+#' @examples
+#' \dontrun{
+#' url <- "https://upload.wikimedia.org/wikipedia/commons/b/b6/Felis_catus-cat_on_snow.jpg"
+#' img <- base_loader(url) |> transform_to_tensor()
+#'
+#' boxes <- torch_tensor(matrix(c(600, 200, 2880, 1860), ncol = 4), dtype = torch_float32())
+#'
+#' before <- list(x = img, y = list(boxes = boxes, labels = "cat"))
+#' class(before) <- c("image_with_bounding_box", "list")
+#'
+#' after <- item_transform_random_resize_crop(before, size = c(600, 800))
+#'
+#' before_plot <- draw_bounding_boxes(before, colors = "blue", width = 10)
+#' after_plot <- draw_bounding_boxes(after, colors = "red", width = 10)
+#' tensor_image_browse(before_plot)
+#' tensor_image_browse(after_plot)
+#' }
+#'
+#' @family item_random_transforms
+#'
+#' @importFrom torch torch_atan2
+#' @export
+item_transform_random_resize_crop <- function(x, size, scale = c(0.08, 1),
+                                              ratio = c(3 / 4, 4 / 3),
+                                              interpolation = 2) {
+  UseMethod("item_transform_random_resize_crop", x)
+}
+
+#' @export
+item_transform_random_resize_crop.dataset <- function(x, size, scale = c(0.08, 1),
+                                                      ratio = c(3 / 4, 4 / 3),
+                                                      interpolation = 2) {
+  original_getitem <- x$.getitem
+  unlockBinding(".getitem", as.environment(x))
+  x$.getitem <- function(index) {
+    item <- original_getitem(index)
+    item_transform_random_resize_crop(item, size = size, scale = scale,
+                                      ratio = ratio, interpolation = interpolation)
+  }
+  x
+}
+
+#' @export
+item_transform_random_resize_crop.default <- function(x, size, scale = c(0.08, 1),
+                                                      ratio = c(3 / 4, 4 / 3),
+                                                      interpolation = 2) {
+  cli_abort(
+    "{.fn item_transform_random_resize_crop} requires a dataset item (a list with {.var x} and {.var y} fields), not {.obj_type_friendly {x}}.
+    To crop and resize a raw image tensor, use {.fn transform_random_resized_crop} instead."
+  )
+}
+
+#' @export
+item_transform_random_resize_crop.image_with_bounding_box <- function(x, size, scale = c(0.08, 1),
+                                                                      ratio = c(3 / 4, 4 / 3),
+                                                                      interpolation = 2) {
+  params <- get_random_resized_crop_params(x$x, scale, ratio)
+  x <- item_transform_crop(x, top = params[1], left = params[2],
+                           height = params[3], width = params[4])
+
+  x$x <- transform_resize(x$x, size, interpolation)
+
+  new_spatial <- tail(x$x$shape, 2)
+  scale_h <- new_spatial[1] / params[3]
+  scale_w <- new_spatial[2] / params[4]
+
+  boxes <- x$y$boxes$clone()
+  if (boxes$size(1) > 0) {
+    boxes[, 1] <- boxes[, 1] * scale_w
+    boxes[, 2] <- boxes[, 2] * scale_h
+    boxes[, 3] <- boxes[, 3] * scale_w
+    boxes[, 4] <- boxes[, 4] * scale_h
+    if (boxes$size(2) == 5) {
+      boxes[, 5] <- rescale_box_angle(boxes[, 5], scale_w, scale_h)
+    }
+  }
+  x$y$boxes <- boxes
+  if (!is.null(x$y$area)) {
+    x$y$area <- x$y$area * (scale_w * scale_h)
+  }
+  x$y$image_height <- new_spatial[1]
+  x$y$image_width <- new_spatial[2]
+
+  x
+}
+
+#' @export
+item_transform_random_resize_crop.image_with_segmentation_mask <- function(x, size, scale = c(0.08, 1),
+                                                                           ratio = c(3 / 4, 4 / 3),
+                                                                           interpolation = 2) {
+  params <- get_random_resized_crop_params(x$x, scale, ratio)
+  x <- item_transform_crop(x, top = params[1], left = params[2],
+                           height = params[3], width = params[4])
+
+  x$x <- transform_resize(x$x, size, interpolation)
+
+  new_spatial <- tail(x$x$shape, 2)
+  masks <- x$y$masks
+  x$y$masks <- if (masks$ndim == 2) {
+    transform_resize(masks$unsqueeze(1), new_spatial, interpolation = 0)$squeeze(1)
+  } else if (masks$size(1) == 0) {
+    torch_zeros(c(0, new_spatial), dtype = masks$dtype, device = masks$device)
+  } else {
+    transform_resize(masks, new_spatial, interpolation = 0)
+  }
+  x$y$image_height <- new_spatial[1]
+  x$y$image_width <- new_spatial[2]
+
+  x
+}
+
+#' @export
+item_transform_random_resize_crop.image_with_rotated_box <- item_transform_random_resize_crop.image_with_bounding_box
+
+rescale_box_angle <- function(angle_deg, scale_w, scale_h) {
+  if (scale_w == scale_h) {
+    return(angle_deg)
+  }
+
+  rad <- deg2rad(angle_deg)
+  rad2deg(torch_atan2(scale_h * torch_sin(rad), scale_w * torch_cos(rad)))
+}

R/transforms-defaults.R

@@ -208,8 +208,8 @@ get_random_resized_crop_params <- function(img, scale, ratio) {
     w <- width
     h <- height
   }
-  i <- (height - h) %/% 2
-  j <- (width - w) %/% 2
+  i <- (height - h) %/% 2 + 1
+  j <- (width - w) %/% 2 + 1
 
   c(i, j, h, w)
 }