Feat: Add roi_align_rotated

NAMESPACE

@@ -2,10 +2,12 @@
 
 export(install_torchvisionlib)
 export(nn_ps_roi_align)
+export(nn_roi_align_rotated)
 export(ops_deform_conv2d)
 export(ops_ms_deform_attn)
 export(ops_nms)
 export(ops_ps_roi_align)
+export(ops_roi_align_rotated)
 export(torchvisionlib_is_installed)
 export(vision_read_jpeg)
 importFrom(Rcpp,sourceCpp)

NEWS.md

@@ -1,5 +1,8 @@
 # torchvisionlib (development version)
 
+- Added `ops_roi_align_rotated()`, a CPU implementation of RoI align pooling
+  for rotated proposals (matching the mmcv `roi_align_rotated` operator). (#32)
+
 - Added `ops_ms_deform_attn()`, a CUDA implementation of multi-scale deformable
   attention (used by Deformable-DETR and LW-DETR). Vendored from Deformable-DETR
   (Apache-2.0). (#25)

R/RcppExports.R

@@ -9,6 +9,10 @@ rcpp_vision_ops_ms_deform_attn <- function(value, spatial_shapes, level_start_in
     .Call('_torchvisionlib_rcpp_vision_ops_ms_deform_attn', PACKAGE = 'torchvisionlib', value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step)
 }
 
+rcpp_vision_ops_roi_align_rotated <- function(input, rois, pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise) {
+    .Call('_torchvisionlib_rcpp_vision_ops_roi_align_rotated', PACKAGE = 'torchvisionlib', input, rois, pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise)
+}
+
 rcpp_vision_ops_deform_conv2d <- function(input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask) {
     .Call('_torchvisionlib_rcpp_vision_ops_deform_conv2d', PACKAGE = 'torchvisionlib', input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask)
 }

R/ops.R

@@ -239,3 +239,79 @@ nn_ps_roi_align <- torch::nn_module(
 )
 
 
+#' RoI align pooling for rotated proposals
+#'
+#' Performs RoI align pooling for rotated proposals, as implemented by the MMCV
+#' `roi_align_rotated` operator
+#' (see <https://mmcv.readthedocs.io/en/latest/deployment/mmcv_ops_definition.html#mmcvroialignrotated>).
+#' Only a CPU implementation is provided.
+#'
+#' @param input (`Tensor[N, C, H, W]`): input feature map.
+#' @param rois (`Tensor[K, 6]`): rotated boxes with columns
+#'   `(batch_index, cx, cy, w, h, angle)`, where `batch_index` is a **0-based**
+#'   index into the first dimension of `input`, `(cx, cy)` is the box center,
+#'   `(w, h)` the box size and `angle` the rotation angle in radians
+#'   (counterclockwise unless `clockwise = TRUE`).
+#' @param output_size (int or `Tuple[int, int]`): the output size `(height, width)`
+#'   after pooling.
+#' @param spatial_scale (float): scaling factor mapping box coordinates to input
+#'   coordinates. For example, if boxes are defined on a 224x224 image and
+#'   `input` is a 112x112 feature map, set this to 0.5.
+#' @param sampling_ratio (int): number of sampling points per output bin. If `<= 0`,
+#'   an adaptive number of grid points is used (computed as `ceil(roi_size / output_size)`).
+#'   Default: 0
+#' @param aligned (bool): if `TRUE` (default), the aligned implementation is
+#'   used (results are shifted by -0.5 before interpolation, matching
+#'   detectron2). If `FALSE`, the legacy MMDetection implementation is used and
+#'   boxes are clamped to a minimum size of 1.
+#' @param clockwise (bool): if `TRUE`, the rotation angle is interpreted in a
+#'   clockwise fashion in image space, otherwise it is counterclockwise.
+#'   Default: `FALSE`
+#'
+#' @returns
+#' `Tensor[K, C, output_size[1], output_size[2]]`: the pooled features, where
+#' the `r`-th element corresponds to the `r`-th RoI in `rois`.
+#'
+#' @examples
+#' if (torchvisionlib_is_installed()) {
+#'   library(torch)
+#'   input <- torch_randn(1, 3, 28, 28)
+#'   # (batch_index, cx, cy, w, h, angle) with 0-based batch index
+#'   rois <- torch_tensor(matrix(c(0, 14, 14, 10, 10, 0.5), ncol = 6))
+#'   ops_roi_align_rotated(input, rois, output_size = c(5, 5),
+#'                         spatial_scale = 1, sampling_ratio = 2)
+#' }
+#'
+#' @family ops
+#' @export
+ops_roi_align_rotated <- function(input, rois, output_size, spatial_scale,
+                                  sampling_ratio = 0, aligned = TRUE,
+                                  clockwise = FALSE) {
+  output_size <- .pair(output_size)
+  rcpp_vision_ops_roi_align_rotated(
+    input, rois,
+    output_size[1], output_size[2],
+    spatial_scale, sampling_ratio,
+    aligned, clockwise
+  )
+}
+
+
+#' @describeIn ops_roi_align_rotated The [torch::nn_module()] wrapper for [ops_roi_align_rotated()].
+#' @export
+nn_roi_align_rotated <- torch::nn_module(
+  initialize = function(output_size, spatial_scale, sampling_ratio = 0,
+                        aligned = TRUE, clockwise = FALSE) {
+    self$output_size <- output_size
+    self$spatial_scale <- spatial_scale
+    self$sampling_ratio <- sampling_ratio
+    self$aligned <- aligned
+    self$clockwise <- clockwise
+  },
+  forward = function(input, rois) {
+    ops_roi_align_rotated(input, rois, self$output_size, self$spatial_scale,
+                          self$sampling_ratio, self$aligned, self$clockwise)
+  }
+)
+
+

csrc/CMakeLists.txt

@@ -149,6 +149,13 @@ if (DEFINED ENV{CUDA} AND NOT '$ENV{CUDA}' STREQUAL '')
   list(APPEND TORCHVISION_SRC ops/ms_deform_attn/cuda/ms_deform_attn_kernel.cu)
 endif()
 
+# RoI align pooling for rotated proposals (mmcv roi_align_rotated, CPU only).
+list(APPEND TORCHVISION_SRC
+    ops/roi_align_rotated/roi_align_rotated.cpp
+    ops/roi_align_rotated/cpu/roi_align_rotated_kernel.cpp
+    ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp
+)
+
 add_library(torchvisionlib SHARED ${TORCHVISION_SRC})
 add_library(torchvisionlib::library ALIAS torchvisionlib)
 

csrc/include/torchvisionlib/exports.h

@@ -29,6 +29,7 @@ TORCHVISIONLIB_API void torchvisionlib_last_error_clear();
 
 TORCHVISIONLIB_API void* _vision_ops_nms (void* dets, void* scores, double iou_threshold);
 TORCHVISIONLIB_API void* _vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void* level_start_index, void* sampling_loc, void* attn_weight, std::int64_t im2col_step);
+TORCHVISIONLIB_API void* _vision_ops_roi_align_rotated (void* input, void* rois, std::int64_t pooled_height, std::int64_t pooled_width, double spatial_scale, std::int64_t sampling_ratio, bool aligned, bool clockwise);
 TORCHVISIONLIB_API void* _vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask);
 TORCHVISIONLIB_API void* _vision_ops_ps_roi_align (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio);
 TORCHVISIONLIB_API void* _vision_ops_ps_roi_pool (void* input, void* rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width);
@@ -51,6 +52,11 @@ inline void* vision_ops_ms_deform_attn (void* value, void* spatial_shapes, void*
   host_exception_handler();
   return ret;
 }
+inline void* vision_ops_roi_align_rotated (void* input, void* rois, std::int64_t pooled_height, std::int64_t pooled_width, double spatial_scale, std::int64_t sampling_ratio, bool aligned, bool clockwise) {
+  auto ret =  _vision_ops_roi_align_rotated(input, rois, pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise);
+  host_exception_handler();
+  return ret;
+}
 inline void* vision_ops_deform_conv2d (void* input, void* weight, void* offset, void* mask, void* bias, std::int64_t stride_h, std::int64_t stride_w, std::int64_t pad_h, std::int64_t pad_w, std::int64_t dilation_h, std::int64_t dilation_w, std::int64_t groups, std::int64_t offset_groups, bool use_mask) {
   auto ret =  _vision_ops_deform_conv2d(input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask);
   host_exception_handler();

csrc/ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp

@@ -0,0 +1,112 @@
+#include "../roi_align_rotated.h"
+
+#include <torch/autograd.h>
+#include <torch/types.h>
+
+namespace vision {
+namespace ops {
+
+namespace {
+
+class ROIAlignRotatedFunction
+    : public torch::autograd::Function<ROIAlignRotatedFunction> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::autograd::Variable& input,
+      const torch::autograd::Variable& rois,
+      int64_t pooled_height,
+      int64_t pooled_width,
+      double spatial_scale,
+      int64_t sampling_ratio,
+      bool aligned,
+      bool clockwise) {
+    at::AutoDispatchBelowADInplaceOrView g;
+    auto output = roi_align_rotated(
+        input,
+        rois,
+        pooled_height,
+        pooled_width,
+        spatial_scale,
+        sampling_ratio,
+        aligned,
+        clockwise);
+
+    ctx->save_for_backward({rois});
+    ctx->saved_data["pooled_height"] = pooled_height;
+    ctx->saved_data["pooled_width"] = pooled_width;
+    ctx->saved_data["spatial_scale"] = spatial_scale;
+    ctx->saved_data["sampling_ratio"] = sampling_ratio;
+    ctx->saved_data["aligned"] = aligned;
+    ctx->saved_data["clockwise"] = clockwise;
+    ctx->saved_data["batch_size"] = input.size(0);
+    ctx->saved_data["channels"] = input.size(1);
+    ctx->saved_data["height"] = input.size(2);
+    ctx->saved_data["width"] = input.size(3);
+
+    return {output};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::autograd::variable_list& grad_output) {
+    auto saved = ctx->get_saved_variables();
+    auto rois = saved[0];
+
+    auto grad_input = detail::_roi_align_rotated_backward(
+        grad_output[0],
+        rois,
+        ctx->saved_data["pooled_height"].toInt(),
+        ctx->saved_data["pooled_width"].toInt(),
+        ctx->saved_data["spatial_scale"].toDouble(),
+        ctx->saved_data["sampling_ratio"].toInt(),
+        ctx->saved_data["aligned"].toBool(),
+        ctx->saved_data["clockwise"].toBool(),
+        ctx->saved_data["batch_size"].toInt(),
+        ctx->saved_data["channels"].toInt(),
+        ctx->saved_data["height"].toInt(),
+        ctx->saved_data["width"].toInt());
+
+    return {
+        grad_input,
+        torch::autograd::Variable(), // rois
+        torch::autograd::Variable(), // pooled_height
+        torch::autograd::Variable(), // pooled_width
+        torch::autograd::Variable(), // spatial_scale
+        torch::autograd::Variable(), // sampling_ratio
+        torch::autograd::Variable(), // aligned
+        torch::autograd::Variable(), // clockwise
+    };
+  }
+};
+
+at::Tensor roi_align_rotated_autograd(
+    const at::Tensor& input,
+    const at::Tensor& rois,
+    int64_t pooled_height,
+    int64_t pooled_width,
+    double spatial_scale,
+    int64_t sampling_ratio,
+    bool aligned,
+    bool clockwise) {
+  return ROIAlignRotatedFunction::apply(
+      input,
+      rois,
+      pooled_height,
+      pooled_width,
+      spatial_scale,
+      sampling_ratio,
+      aligned,
+      clockwise)[0];
+}
+
+} // namespace
+
+TORCH_LIBRARY_IMPL(torchvision, Autograd, m) {
+  m.impl(
+      TORCH_SELECTIVE_NAME("torchvision::roi_align_rotated"),
+      TORCH_FN(roi_align_rotated_autograd));
+}
+
+} // namespace ops
+} // namespace vision