Feat: item transform random affine

NAMESPACE

@@ -29,6 +29,11 @@ S3method(item_transform_hflip,default)
 S3method(item_transform_hflip,image_with_bounding_box)
 S3method(item_transform_hflip,image_with_rotated_box)
 S3method(item_transform_hflip,image_with_segmentation_mask)
+S3method(item_transform_random_affine,dataset)
+S3method(item_transform_random_affine,default)
+S3method(item_transform_random_affine,image_with_bounding_box)
+S3method(item_transform_random_affine,image_with_rotated_box)
+S3method(item_transform_random_affine,image_with_segmentation_mask)
 S3method(item_transform_random_horizontal_flip,dataset)
 S3method(item_transform_random_horizontal_flip,default)
 S3method(item_transform_random_horizontal_flip,image_with_bounding_box)
@@ -171,6 +176,7 @@ export(item_transform_affine)
 export(item_transform_center_crop)
 export(item_transform_crop)
 export(item_transform_hflip)
+export(item_transform_random_affine)
 export(item_transform_random_horizontal_flip)
 export(item_transform_random_vertical_flip)
 export(item_transform_pad)

NEWS.md

@@ -26,9 +26,12 @@
 * `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_affine()` for applying an affine transformation drawn from the given ranges to dataset items, with support for detection, segmentation and rotated-box item types and datasets (@srishtiii28, #354).
 
 ## Bug fixes and improvements
 
+* `transform_random_affine()` now accepts a bare number for `shear`. It used to widen `degrees`
+  instead of `shear`, which left the shear range incomplete and made the sampling fail (#354).
 * `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,115 @@ item_transform_random_vertical_flip.image_with_rotated_box <- function(x, p = 0.
   }
   x
 }
+
+#' Randomly apply an affine transformation on a dataset item
+#'
+#' Draws a random rotation, translation, scale and shear inside the given ranges
+#' and applies the resulting affine transformation to the dataset item with
+#' \code{\link{item_transform_affine}}, keeping the image size unchanged. Image
+#' and target share the same draw, so that they stay aligned.
+#'
+#' The transformation is drawn again for every item, so that a dataset wrapped
+#' with this transform yields a different transformation on each access.
+#'
+#' @param x A dataset item, typically an \code{image_with_bounding_box},
+#'   \code{image_with_rotated_box} or \code{image_with_segmentation_mask} object
+#'   containing an image tensor and associated target data.
+#' @inheritParams transform_random_affine
+#' @param fill Fill color for the area outside the transform. Default is
+#'   \code{NULL}.
+#' @param center (numeric vector of length 2, optional): Optional center of
+#'   rotation, \code{c(x, y)}. Default is image center.
+#'
+#' @return A dataset item with the image and target transformed. Detection items
+#'   are returned as \code{image_with_rotated_box}; segmentation items keep their
+#'   class.
+#'
+#' @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",
+#'                                  image_height = img$shape[2], image_width = img$shape[3]))
+#' class(before) <- c("image_with_bounding_box", "list")
+#'
+#' after <- item_transform_random_affine(before, degrees = 30, translate = c(0.1, 0.1),
+#'                                       scale = c(0.8, 1.2), shear = 10)
+#'
+#' 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
+#'
+#' @export
+item_transform_random_affine <- function(x, degrees, translate = NULL, scale = NULL,
+                                         shear = NULL, interpolation = 0, fill = NULL,
+                                         center = NULL) {
+  UseMethod("item_transform_random_affine", x)
+}
+
+#' @export
+item_transform_random_affine.default <- function(x, degrees, translate = NULL, scale = NULL,
+                                                 shear = NULL, interpolation = 0, fill = NULL,
+                                                 center = NULL) {
+  cli_abort(
+    "{.fn item_transform_random_affine} requires a dataset item (a list with {.var x} and {.var y} fields), not {.obj_type_friendly {x}}.
+    To transform a raw image tensor, use {.fn transform_random_affine} instead."
+  )
+}
+
+#' @export
+item_transform_random_affine.dataset <- function(x, degrees, translate = NULL, scale = NULL,
+                                                 shear = NULL, interpolation = 0, fill = NULL,
+                                                 center = NULL) {
+  original_getitem <- x$.getitem
+  unlockBinding(".getitem", as.environment(x))
+  x$.getitem <- function(index) {
+    item <- original_getitem(index)
+    item_transform_random_affine(item, degrees = degrees, translate = translate,
+                                 scale = scale, shear = shear,
+                                 interpolation = interpolation, fill = fill,
+                                 center = center)
+  }
+  x
+}
+
+#' @export
+item_transform_random_affine.image_with_bounding_box <- function(x, degrees, translate = NULL,
+                                                                 scale = NULL, shear = NULL,
+                                                                 interpolation = 0, fill = NULL,
+                                                                 center = NULL) {
+  random_affine_item(x, degrees, translate, scale, shear, interpolation, fill, center)
+}
+
+#' @export
+item_transform_random_affine.image_with_rotated_box <- function(x, degrees, translate = NULL,
+                                                                scale = NULL, shear = NULL,
+                                                                interpolation = 0, fill = NULL,
+                                                                center = NULL) {
+  random_affine_item(x, degrees, translate, scale, shear, interpolation, fill, center)
+}
+
+#' @export
+item_transform_random_affine.image_with_segmentation_mask <- function(x, degrees, translate = NULL,
+                                                                      scale = NULL, shear = NULL,
+                                                                      interpolation = 0, fill = NULL,
+                                                                      center = NULL) {
+  random_affine_item(x, degrees, translate, scale, shear, interpolation, fill, center)
+}
+
+random_affine_item <- function(x, degrees, translate, scale, shear, interpolation, fill, center) {
+  args <- check_random_affine_params(degrees, translate, scale, shear)
+  params <- get_random_affine_params(args$degrees, translate, scale, args$shear,
+                                     get_image_size(x$x))
+
+  item_transform_affine(x, angle = params[[1]], translate = params[[2]],
+                        scale = params[[3]], shear = params[[4]],
+                        interpolation = interpolation, fill = fill, center = center)
+}

R/transforms-defaults.R

@@ -366,6 +366,57 @@ transform_random_rotation.default <- function(img, degrees, interpolation=0,
 }
 
 
+check_random_affine_params <- function(degrees, translate, scale, shear) {
+  if (length(degrees) == 1) {
+
+    if (degrees < 0)
+      value_error("degrees must be positive if it's a single value")
+
+    degrees <- c(-degrees, degrees)
+
+  } else if (length(degrees) != 2) {
+    value_error("degrees must be length 1 or 2")
+  }
+
+
+  if (!is.null(translate)) {
+
+    if (length(translate) != 2)
+      value_error("translate must be length 2")
+
+    if (any(translate > 1) || any(translate < 0))
+      value_error("translate must be between 0 and 1")
+
+  }
+
+  if (!is.null(scale)) {
+
+    if (length(scale) != 2)
+      value_error("scale must be length 2")
+
+    if (any(scale < 0))
+      value_error("scale must be positive")
+
+  }
+
+  if (!is.null(shear)) {
+
+    if (length(shear) == 1) {
+
+      if (shear < 0)
+        value_error("shear must be positive if it's a single value")
+
+      shear <- c(-shear, shear)
+
+    } else if (!length(shear) %in% c(2, 4)) {
+      value_error("shear's length must be 1, 2, or 4")
+    }
+
+  }
+
+  list(degrees = degrees, shear = shear)
+}
+
 get_random_affine_params <- function(degrees,
                                      translate,
                                      scale_ranges,
@@ -415,57 +466,11 @@ transform_random_affine.default <- function(img, degrees, translate=NULL, scale=
     fill <- fillcolor
   }
 
-  if (length(degrees) == 1) {
-
-    if (degrees < 0)
-      value_error("degrees must be positive if it's a single value")
-
-    degrees <- c(-degrees, degrees)
-
-  } else if (length(degrees) != 2) {
-    value_error("degrees must be length 1 or 2")
-  }
-
-
-  if (!is.null(translate)) {
-
-    if (length(translate) != 2)
-      value_error("translate must be length 2")
-
-    if (any(translate > 1) || any(translate < 0))
-      value_error("translate must be between 0 and 1")
-
-  }
-
-  if (!is.null(scale)) {
-
-    if (length(scale) != 2)
-      value_error("scale must be length 2")
-
-    if (any(scale < 0))
-      value_error("scale must be positive")
-
-  }
-
-  if (!is.null(shear)) {
-
-    if (length(shear) == 1) {
-
-      if (shear < 0)
-        value_error("shear must be positive if it's a single value")
-
-      degrees <- c(-degrees, degrees)
-
-    } else if (!length(shear) %in% c(2, 4)) {
-      value_error("shear's length must be 1, 2, or 4")
-    }
-
-  }
-
+  args <- check_random_affine_params(degrees, translate, scale, shear)
 
   img_size <- get_image_size(img)
 
-  ret <- get_random_affine_params(degrees, translate, scale, shear, img_size)
+  ret <- get_random_affine_params(args$degrees, translate, scale, args$shear, img_size)
 
   transform_affine(img, ret[[1]], ret[[2]], ret[[3]], ret[[4]],
                    interpolation=interpolation, fill=fill)