From 7eb9e393b2c007400933f572a245ffe770453750 Mon Sep 17 00:00:00 2001 From: DerrickUnleashed Date: Mon, 3 Aug 2026 00:02:43 +0530 Subject: [PATCH 1/4] RoI align rotated CPU op (#32) --- NAMESPACE | 2 + NEWS.md | 3 + R/RcppExports.R | 4 + R/ops.R | 76 +++ csrc/CMakeLists.txt | 7 + csrc/include/torchvisionlib/exports.h | 6 + .../autograd/roi_align_rotated_kernel.cpp | 112 ++++ .../cpu/roi_align_rotated_kernel.cpp | 575 ++++++++++++++++++ .../roi_align_rotated/roi_align_rotated.cpp | 77 +++ .../ops/roi_align_rotated/roi_align_rotated.h | 44 ++ csrc/src/exports.cpp | 7 + csrc/src/ops.cpp | 23 + csrc/src/torchvisionlib.def | 1 + inst/def/torchvisionlib.def | 1 + man/ops_ms_deform_attn.Rd | 3 +- man/ops_nms.Rd | 3 +- man/ops_roi_align_rotated.Rd | 86 +++ src/RcppExports.cpp | 19 + src/exports.cpp | 4 + src/exports.h | 1 + tests/run-roi-align-rotated.R | 197 ++++++ tests/testthat/test-ops-roi-align-rotated.R | 205 +++++++ 22 files changed, 1454 insertions(+), 2 deletions(-) create mode 100644 csrc/ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp create mode 100644 csrc/ops/roi_align_rotated/cpu/roi_align_rotated_kernel.cpp create mode 100644 csrc/ops/roi_align_rotated/roi_align_rotated.cpp create mode 100644 csrc/ops/roi_align_rotated/roi_align_rotated.h create mode 100644 man/ops_roi_align_rotated.Rd create mode 100644 tests/run-roi-align-rotated.R create mode 100644 tests/testthat/test-ops-roi-align-rotated.R diff --git a/NAMESPACE b/NAMESPACE index 064fb6e..69ebac3 100644 --- a/NAMESPACE +++ b/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) diff --git a/NEWS.md b/NEWS.md index ec78e6b..fe6ed59 100644 --- a/NEWS.md +++ b/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) diff --git a/R/RcppExports.R b/R/RcppExports.R index 46e8c79..e04fad5 100644 --- a/R/RcppExports.R +++ b/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) } diff --git a/R/ops.R b/R/ops.R index c6920b5..d851474 100644 --- a/R/ops.R +++ b/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 ). +#' 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) + } +) + + diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 9cdf0d9..61622e0 100644 --- a/csrc/CMakeLists.txt +++ b/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) diff --git a/csrc/include/torchvisionlib/exports.h b/csrc/include/torchvisionlib/exports.h index 16278a9..b991ac8 100644 --- a/csrc/include/torchvisionlib/exports.h +++ b/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(); diff --git a/csrc/ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp b/csrc/ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp new file mode 100644 index 0000000..9c5b6cb --- /dev/null +++ b/csrc/ops/roi_align_rotated/autograd/roi_align_rotated_kernel.cpp @@ -0,0 +1,112 @@ +#include "../roi_align_rotated.h" + +#include +#include + +namespace vision { +namespace ops { + +namespace { + +class ROIAlignRotatedFunction + : public torch::autograd::Function { + 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 diff --git a/csrc/ops/roi_align_rotated/cpu/roi_align_rotated_kernel.cpp b/csrc/ops/roi_align_rotated/cpu/roi_align_rotated_kernel.cpp new file mode 100644 index 0000000..f5f8479 --- /dev/null +++ b/csrc/ops/roi_align_rotated/cpu/roi_align_rotated_kernel.cpp @@ -0,0 +1,575 @@ +// Modified from MMCV: +// https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/pytorch/cpu/roi_align_rotated_cpu.cpp +// which is itself a port of detectron2's `ROIAlignRotated`: +// https://github.com/facebookresearch/detectron2/tree/master/detectron2/layers/csrc/ROIAlignRotated +// Copyright (c) OpenMMLab / Facebook, Inc. and its affiliates. All Rights Reserved +#include +#include +#include + +#include +#include + +namespace vision { +namespace ops { + +namespace { + +template +struct PreCalc { + int pos1; + int pos2; + int pos3; + int pos4; + T w1; + T w2; + T w3; + T w4; +}; + +template +void pre_calc_for_bilinear_interpolate( + const int height, + const int width, + const int pooled_height, + const int pooled_width, + const int iy_upper, + const int ix_upper, + T roi_start_h, + T roi_start_w, + T bin_size_h, + T bin_size_w, + int roi_bin_grid_h, + int roi_bin_grid_w, + T roi_center_h, + T roi_center_w, + T cos_theta, + T sin_theta, + std::vector>& pre_calc) { + int pre_calc_index = 0; + for (int ph = 0; ph < pooled_height; ph++) { + for (int pw = 0; pw < pooled_width; pw++) { + for (int iy = 0; iy < iy_upper; iy++) { + const T yy = roi_start_h + ph * bin_size_h + + static_cast(iy + .5f) * bin_size_h / + static_cast(roi_bin_grid_h); // e.g., 0.5, 1.5 + for (int ix = 0; ix < ix_upper; ix++) { + const T xx = roi_start_w + pw * bin_size_w + + static_cast(ix + .5f) * bin_size_w / + static_cast(roi_bin_grid_w); + + // Rotate by theta around the center and translate + // In image space, (y, x) is the order for Right Handed System, + // and this is essentially multiplying the point by a rotation matrix + // to rotate it counterclockwise through angle theta. + T y = yy * cos_theta - xx * sin_theta + roi_center_h; + T x = yy * sin_theta + xx * cos_theta + roi_center_w; + // deal with: inverse elements are out of feature map boundary + if (y < -1.0 || y > height || x < -1.0 || x > width) { + // empty + PreCalc pc; + pc.pos1 = 0; + pc.pos2 = 0; + pc.pos3 = 0; + pc.pos4 = 0; + pc.w1 = 0; + pc.w2 = 0; + pc.w3 = 0; + pc.w4 = 0; + pre_calc[pre_calc_index] = pc; + pre_calc_index += 1; + continue; + } + + if (y < 0) { + y = 0; + } + if (x < 0) { + x = 0; + } + + int y_low = (int)y; + int x_low = (int)x; + int y_high; + int x_high; + + if (y_low >= height - 1) { + y_high = y_low = height - 1; + y = (T)y_low; + } else { + y_high = y_low + 1; + } + + if (x_low >= width - 1) { + x_high = x_low = width - 1; + x = (T)x_low; + } else { + x_high = x_low + 1; + } + + T ly = y - y_low; + T lx = x - x_low; + T hy = 1. - ly, hx = 1. - lx; + T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; + + // save weights and indices + PreCalc pc; + pc.pos1 = y_low * width + x_low; + pc.pos2 = y_low * width + x_high; + pc.pos3 = y_high * width + x_low; + pc.pos4 = y_high * width + x_high; + pc.w1 = w1; + pc.w2 = w2; + pc.w3 = w3; + pc.w4 = w4; + pre_calc[pre_calc_index] = pc; + + pre_calc_index += 1; + } + } + } + } +} + +template +void ROIAlignRotatedForward( + const int nthreads, + const T* input, + const T& spatial_scale, + const bool aligned, + const bool clockwise, + const int channels, + const int height, + const int width, + const int pooled_height, + const int pooled_width, + const int sampling_ratio, + const T* rois, + T* output) { + int n_rois = nthreads / channels / pooled_width / pooled_height; + // (n, c, ph, pw) is an element in the pooled output. + // Each ROI writes to a distinct region of `output` and only reads `input`, + // so the ROI dimension can be safely parallelized. + at::parallel_for(0, n_rois, 1, [&](int64_t start, int64_t end) { + for (int64_t n = start; n < end; n++) { + int64_t index_n = n * channels * pooled_width * pooled_height; + + const T* current_roi = rois + n * 6; + int roi_batch_ind = current_roi[0]; + + // Do not use rounding; this implementation detail is critical + T offset = aligned ? (T)0.5 : (T)0.0; + T roi_center_w = current_roi[1] * spatial_scale - offset; + T roi_center_h = current_roi[2] * spatial_scale - offset; + T roi_width = current_roi[3] * spatial_scale; + T roi_height = current_roi[4] * spatial_scale; + T theta = current_roi[5]; + if (clockwise) { + theta = -theta; // If clockwise, the angle needs to be reversed. + } + T cos_theta = cos(theta); + T sin_theta = sin(theta); + + if (aligned) { + TORCH_CHECK( + roi_width >= 0 && roi_height >= 0, + "ROIs in ROIAlignRotated do not have non-negative size!"); + } else { // for backward-compatibility only + roi_width = std::max(roi_width, (T)1.); + roi_height = std::max(roi_height, (T)1.); + } + + T bin_size_h = static_cast(roi_height) / static_cast(pooled_height); + T bin_size_w = static_cast(roi_width) / static_cast(pooled_width); + + // We use roi_bin_grid to sample the grid and mimic integral + int roi_bin_grid_h = (sampling_ratio > 0) + ? sampling_ratio + : ceilf(roi_height / pooled_height); // e.g., = 2 + int roi_bin_grid_w = + (sampling_ratio > 0) ? sampling_ratio : ceilf(roi_width / pooled_width); + + // We do average (integral) pooling inside a bin + const T count = std::max(roi_bin_grid_h * roi_bin_grid_w, 1); // e.g. = 4 + + // we want to precalculate indices and weights shared by all channels, + // this is the key point of optimization + std::vector> pre_calc(roi_bin_grid_h * roi_bin_grid_w * + pooled_width * pooled_height); + + // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y). + // Appropriate translation needs to be applied after. + T roi_start_h = -roi_height / 2.0; + T roi_start_w = -roi_width / 2.0; + + pre_calc_for_bilinear_interpolate( + height, width, pooled_height, pooled_width, roi_bin_grid_h, + roi_bin_grid_w, roi_start_h, roi_start_w, bin_size_h, bin_size_w, + roi_bin_grid_h, roi_bin_grid_w, roi_center_h, roi_center_w, cos_theta, + sin_theta, pre_calc); + + for (int c = 0; c < channels; c++) { + int64_t index_n_c = index_n + c * pooled_width * pooled_height; + const T* offset_input = + input + (roi_batch_ind * channels + c) * height * width; + int pre_calc_index = 0; + + for (int ph = 0; ph < pooled_height; ph++) { + for (int pw = 0; pw < pooled_width; pw++) { + int64_t index = index_n_c + ph * pooled_width + pw; + + T output_val = 0.; + for (int iy = 0; iy < roi_bin_grid_h; iy++) { + for (int ix = 0; ix < roi_bin_grid_w; ix++) { + PreCalc pc = pre_calc[pre_calc_index]; + output_val += pc.w1 * offset_input[pc.pos1] + + pc.w2 * offset_input[pc.pos2] + + pc.w3 * offset_input[pc.pos3] + + pc.w4 * offset_input[pc.pos4]; + + pre_calc_index += 1; + } + } + output_val /= count; + + output[index] = output_val; + } // for pw + } // for ph + } // for c + } // for n + }); +} + +template +void bilinear_interpolate_gradient( + const int height, + const int width, + T y, + T x, + T& w1, + T& w2, + T& w3, + T& w4, + int& x_low, + int& x_high, + int& y_low, + int& y_high) { + // deal with cases that inverse elements are out of feature map boundary + if (y < -1.0 || y > height || x < -1.0 || x > width) { + // empty + w1 = w2 = w3 = w4 = 0.; + x_low = x_high = y_low = y_high = -1; + return; + } + + if (y < 0) { + y = 0; + } + + if (x < 0) { + x = 0; + } + + y_low = (int)y; + x_low = (int)x; + + if (y_low >= height - 1) { + y_high = y_low = height - 1; + y = (T)y_low; + } else { + y_high = y_low + 1; + } + + if (x_low >= width - 1) { + x_high = x_low = width - 1; + x = (T)x_low; + } else { + x_high = x_low + 1; + } + + T ly = y - y_low; + T lx = x - x_low; + T hy = 1. - ly, hx = 1. - lx; + + w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; +} + +template +inline void add(T* address, const T& val) { + *address += val; +} + +template +void ROIAlignRotatedBackward( + const int nthreads, + // may not be contiguous. should index using n_stride, etc + const T* grad_output, + const T& spatial_scale, + const bool aligned, + const bool clockwise, + const int channels, + const int height, + const int width, + const int pooled_height, + const int pooled_width, + const int sampling_ratio, + T* grad_input, + const T* rois, + const int n_stride, + const int c_stride, + const int h_stride, + const int w_stride) { + // Sampling regions of different bins within a single ROI can overlap in + // input space, so this kernel must run serially to accumulate gradients + // without races. + for (int index = 0; index < nthreads; index++) { + // (n, c, ph, pw) is an element in the pooled output + int pw = index % pooled_width; + int ph = (index / pooled_width) % pooled_height; + int c = (index / pooled_width / pooled_height) % channels; + int n = index / pooled_width / pooled_height / channels; + + const T* current_roi = rois + n * 6; + int roi_batch_ind = current_roi[0]; + + // Do not use rounding; this implementation detail is critical + T offset = aligned ? (T)0.5 : (T)0.0; + T roi_center_w = current_roi[1] * spatial_scale - offset; + T roi_center_h = current_roi[2] * spatial_scale - offset; + T roi_width = current_roi[3] * spatial_scale; + T roi_height = current_roi[4] * spatial_scale; + T theta = current_roi[5]; + if (clockwise) { + theta = -theta; // If clockwise, the angle needs to be reversed. + } + T cos_theta = cos(theta); + T sin_theta = sin(theta); + + if (aligned) { + TORCH_CHECK( + roi_width >= 0 && roi_height >= 0, + "ROIs in ROIAlignRotated do not have non-negative size!"); + } else { // for backward-compatibility only + roi_width = std::max(roi_width, (T)1.); + roi_height = std::max(roi_height, (T)1.); + } + + T bin_size_h = static_cast(roi_height) / static_cast(pooled_height); + T bin_size_w = static_cast(roi_width) / static_cast(pooled_width); + + T* offset_grad_input = + grad_input + ((roi_batch_ind * channels + c) * height * width); + + int output_offset = n * n_stride + c * c_stride; + const T* offset_grad_output = grad_output + output_offset; + const T grad_output_this_bin = + offset_grad_output[ph * h_stride + pw * w_stride]; + + // We use roi_bin_grid to sample the grid and mimic integral + int roi_bin_grid_h = (sampling_ratio > 0) + ? sampling_ratio + : ceilf(roi_height / pooled_height); // e.g., = 2 + int roi_bin_grid_w = + (sampling_ratio > 0) ? sampling_ratio : ceilf(roi_width / pooled_width); + + // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y). + // Appropriate translation needs to be applied after. + T roi_start_h = -roi_height / 2.0; + T roi_start_w = -roi_width / 2.0; + + // We do average (integral) pooling inside a bin + const T count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4 + + for (int iy = 0; iy < roi_bin_grid_h; iy++) { + const T yy = roi_start_h + ph * bin_size_h + + static_cast(iy + .5f) * bin_size_h / + static_cast(roi_bin_grid_h); // e.g., 0.5, 1.5 + for (int ix = 0; ix < roi_bin_grid_w; ix++) { + const T xx = roi_start_w + pw * bin_size_w + + static_cast(ix + .5f) * bin_size_w / + static_cast(roi_bin_grid_w); + + // Rotate by theta around the center and translate + T y = yy * cos_theta - xx * sin_theta + roi_center_h; + T x = yy * sin_theta + xx * cos_theta + roi_center_w; + + T w1, w2, w3, w4; + int x_low, x_high, y_low, y_high; + + bilinear_interpolate_gradient(height, width, y, x, w1, w2, w3, w4, + x_low, x_high, y_low, y_high); + + T g1 = grad_output_this_bin * w1 / count; + T g2 = grad_output_this_bin * w2 / count; + T g3 = grad_output_this_bin * w3 / count; + T g4 = grad_output_this_bin * w4 / count; + + if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) { + add(offset_grad_input + y_low * width + x_low, static_cast(g1)); + add(offset_grad_input + y_low * width + x_high, static_cast(g2)); + add(offset_grad_input + y_high * width + x_low, static_cast(g3)); + add(offset_grad_input + y_high * width + x_high, static_cast(g4)); + } // if + } // ix + } // iy + } // for +} + +void check_rois( + const at::Tensor& input, + const at::Tensor& rois, + const int64_t pooled_height, + const int64_t pooled_width) { + TORCH_CHECK(input.dim() == 4, "input should be a 4d tensor, got ", input.dim(), "D"); + TORCH_CHECK(rois.dim() == 2, "rois should be a 2d tensor, got ", rois.dim(), "D"); + TORCH_CHECK( + rois.size(1) == 6, + "rois should have 6 columns (batch_index, cx, cy, w, h, angle), got ", + rois.size(1)); + TORCH_CHECK( + pooled_height > 0 && pooled_width > 0, + "pooled_height and pooled_width should be positive, got ", + pooled_height, + " and ", + pooled_width); + + const int64_t batch_size = input.size(0); + const int64_t num_rois = rois.size(0); + auto rois_double = rois.to(at::kDouble).contiguous(); + const double* rois_ptr = rois_double.data_ptr(); + for (int64_t i = 0; i < num_rois; i++) { + const double roi_batch_ind = rois_ptr[i * 6]; + TORCH_CHECK( + roi_batch_ind >= 0 && roi_batch_ind < batch_size, + "rois index should be in [0, batch_size), got ", + roi_batch_ind, + " in batch of size ", + batch_size); + } +} + +at::Tensor roi_align_rotated_forward_kernel( + 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) { + TORCH_CHECK(input.is_cpu(), "roi_align_rotated is only implemented on the CPU"); + TORCH_CHECK(rois.is_cpu(), "rois must be a CPU tensor"); + TORCH_CHECK( + rois.scalar_type() == input.scalar_type(), + "rois should have the same dtype as input, got ", + rois.scalar_type(), + " and ", + input.scalar_type()); + check_rois(input, rois, pooled_height, pooled_width); + + const int64_t channels = input.size(1); + const int64_t height = input.size(2); + const int64_t width = input.size(3); + const int64_t num_rois = rois.size(0); + + auto output = at::empty( + {num_rois, channels, pooled_height, pooled_width}, input.options()); + + auto input_contig = input.contiguous(); + auto rois_contig = rois.contiguous(); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "roi_align_rotated_forward", [&] { + ROIAlignRotatedForward( + output.numel(), + input_contig.data_ptr(), + static_cast(spatial_scale), + aligned, + clockwise, + static_cast(channels), + static_cast(height), + static_cast(width), + static_cast(pooled_height), + static_cast(pooled_width), + static_cast(sampling_ratio), + rois_contig.data_ptr(), + output.data_ptr()); + }); + + return output; +} + +at::Tensor roi_align_rotated_backward_kernel( + const at::Tensor& grad_output, + const at::Tensor& rois, + int64_t pooled_height, + int64_t pooled_width, + double spatial_scale, + int64_t sampling_ratio, + bool aligned, + bool clockwise, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width) { + TORCH_CHECK(grad_output.is_cpu(), "roi_align_rotated is only implemented on the CPU"); + TORCH_CHECK(rois.is_cpu(), "rois must be a CPU tensor"); + TORCH_CHECK( + rois.scalar_type() == grad_output.scalar_type(), + "rois should have the same dtype as grad_output, got ", + rois.scalar_type(), + " and ", + grad_output.scalar_type()); + TORCH_CHECK( + batch_size > 0 && channels > 0 && height > 0 && width > 0, + "input feature shape must be positive, got ", + batch_size, + ", ", + channels, + ", ", + height, + ", ", + width); + + auto grad_input = at::zeros( + {batch_size, channels, height, width}, grad_output.options()); + + auto grad_output_contig = grad_output.contiguous(); + auto rois_contig = rois.contiguous(); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_output.scalar_type(), "roi_align_rotated_backward", [&] { + ROIAlignRotatedBackward( + grad_output_contig.numel(), + grad_output_contig.data_ptr(), + static_cast(spatial_scale), + aligned, + clockwise, + static_cast(channels), + static_cast(height), + static_cast(width), + static_cast(pooled_height), + static_cast(pooled_width), + static_cast(sampling_ratio), + grad_input.data_ptr(), + rois_contig.data_ptr(), + static_cast(grad_output_contig.stride(0)), + static_cast(grad_output_contig.stride(1)), + static_cast(grad_output_contig.stride(2)), + static_cast(grad_output_contig.stride(3))); + }); + + return grad_input; +} + +} // namespace + +TORCH_LIBRARY_IMPL(torchvision, CPU, m) { + m.impl( + TORCH_SELECTIVE_NAME("torchvision::roi_align_rotated"), + TORCH_FN(roi_align_rotated_forward_kernel)); + m.impl( + TORCH_SELECTIVE_NAME("torchvision::_roi_align_rotated_backward"), + TORCH_FN(roi_align_rotated_backward_kernel)); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/roi_align_rotated/roi_align_rotated.cpp b/csrc/ops/roi_align_rotated/roi_align_rotated.cpp new file mode 100644 index 0000000..e46abf3 --- /dev/null +++ b/csrc/ops/roi_align_rotated/roi_align_rotated.cpp @@ -0,0 +1,77 @@ +#include "roi_align_rotated.h" + +#include +#include +#include + +namespace vision { +namespace ops { + +at::Tensor roi_align_rotated( + 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) { + static auto op = c10::Dispatcher::singleton() + .findSchemaOrThrow("torchvision::roi_align_rotated", "") + .typed(); + return op.call( + input, + rois, + pooled_height, + pooled_width, + spatial_scale, + sampling_ratio, + aligned, + clockwise); +} + +namespace detail { + +at::Tensor _roi_align_rotated_backward( + const at::Tensor& grad_output, + const at::Tensor& rois, + int64_t pooled_height, + int64_t pooled_width, + double spatial_scale, + int64_t sampling_ratio, + bool aligned, + bool clockwise, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width) { + static auto op = + c10::Dispatcher::singleton() + .findSchemaOrThrow("torchvision::_roi_align_rotated_backward", "") + .typed(); + return op.call( + grad_output, + rois, + pooled_height, + pooled_width, + spatial_scale, + sampling_ratio, + aligned, + clockwise, + batch_size, + channels, + height, + width); +} + +} // namespace detail + +TORCH_LIBRARY_FRAGMENT(torchvision, m) { + m.def(TORCH_SELECTIVE_SCHEMA( + "torchvision::roi_align_rotated(Tensor input, Tensor rois, int pooled_height, int pooled_width, float spatial_scale, int sampling_ratio, bool aligned, bool clockwise) -> Tensor")); + m.def(TORCH_SELECTIVE_SCHEMA( + "torchvision::_roi_align_rotated_backward(Tensor grad_output, Tensor rois, int pooled_height, int pooled_width, float spatial_scale, int sampling_ratio, bool aligned, bool clockwise, int batch_size, int channels, int height, int width) -> Tensor")); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/roi_align_rotated/roi_align_rotated.h b/csrc/ops/roi_align_rotated/roi_align_rotated.h new file mode 100644 index 0000000..73a27d7 --- /dev/null +++ b/csrc/ops/roi_align_rotated/roi_align_rotated.h @@ -0,0 +1,44 @@ +#pragma once + +#include + +namespace vision { +namespace ops { + +// RoI align pooling for rotated proposals, matching the MMCV +// `roi_align_rotated` operator: +// https://mmcv.readthedocs.io/en/latest/deployment/mmcv_ops_definition.html#mmcvroialignrotated +// +// rois is a `Tensor[K, 6]` with columns `(batch_index, cx, cy, w, h, angle)` +// where `batch_index` is 0-based, the box is centered at `(cx, cy)` with size +// `(w, h)` and `angle` is expressed in radians. +at::Tensor roi_align_rotated( + 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); + +namespace detail { + +at::Tensor _roi_align_rotated_backward( + const at::Tensor& grad_output, + const at::Tensor& rois, + int64_t pooled_height, + int64_t pooled_width, + double spatial_scale, + int64_t sampling_ratio, + bool aligned, + bool clockwise, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width); + +} // namespace detail + +} // namespace ops +} // namespace vision diff --git a/csrc/src/exports.cpp b/csrc/src/exports.cpp index 94bfc5f..e1915d9 100644 --- a/csrc/src/exports.cpp +++ b/csrc/src/exports.cpp @@ -28,6 +28,13 @@ TORCHVISIONLIB_API void* _vision_ops_ms_deform_attn (void* value, void* spatial_ } TORCHVISIONLIB_HANDLE_EXCEPTION return (void*) NULL; } +torch::Tensor vision_ops_roi_align_rotated (torch::Tensor input, torch::Tensor 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_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) { + try { + return make_raw::Tensor(vision_ops_roi_align_rotated(from_raw::Tensor(input), from_raw::Tensor(rois), pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise)); + } TORCHVISIONLIB_HANDLE_EXCEPTION + return (void*) NULL; +} torch::Tensor vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor 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_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) { try { diff --git a/csrc/src/ops.cpp b/csrc/src/ops.cpp index 74f3eab..1ffeef3 100644 --- a/csrc/src/ops.cpp +++ b/csrc/src/ops.cpp @@ -8,6 +8,7 @@ #include #include #include "ops/ms_deform_attn/ms_deform_attn.h" +#include "ops/roi_align_rotated/roi_align_rotated.h" // [[torch::export]] torch::Tensor vision_ops_nms(torch::Tensor dets, torch::Tensor scores, double iou_threshold) { @@ -32,6 +33,28 @@ torch::Tensor vision_ops_ms_deform_attn( ); } +// [[torch::export]] +torch::Tensor vision_ops_roi_align_rotated( + torch::Tensor input, + torch::Tensor rois, + std::int64_t pooled_height, + std::int64_t pooled_width, + double spatial_scale, + std::int64_t sampling_ratio, + bool aligned, + bool clockwise) { + return vision::ops::roi_align_rotated( + input, + rois, + pooled_height, + pooled_width, + spatial_scale, + sampling_ratio, + aligned, + clockwise + ); +} + // [[torch::export]] torch::Tensor vision_ops_deform_conv2d( torch::Tensor input, diff --git a/csrc/src/torchvisionlib.def b/csrc/src/torchvisionlib.def index 198d070..490e202 100644 --- a/csrc/src/torchvisionlib.def +++ b/csrc/src/torchvisionlib.def @@ -5,6 +5,7 @@ EXPORTS ; don't modify between the autogenerated lines _vision_ops_nms _vision_ops_ms_deform_attn + _vision_ops_roi_align_rotated _vision_ops_deform_conv2d _vision_ops_ps_roi_align _vision_ops_ps_roi_pool diff --git a/inst/def/torchvisionlib.def b/inst/def/torchvisionlib.def index 198d070..490e202 100644 --- a/inst/def/torchvisionlib.def +++ b/inst/def/torchvisionlib.def @@ -5,6 +5,7 @@ EXPORTS ; don't modify between the autogenerated lines _vision_ops_nms _vision_ops_ms_deform_attn + _vision_ops_roi_align_rotated _vision_ops_deform_conv2d _vision_ops_ps_roi_align _vision_ops_ps_roi_pool diff --git a/man/ops_ms_deform_attn.Rd b/man/ops_ms_deform_attn.Rd index af34c87..b9b4340 100644 --- a/man/ops_ms_deform_attn.Rd +++ b/man/ops_ms_deform_attn.Rd @@ -49,6 +49,7 @@ feature tensor, so no 1-based index adjustment is applied. } \seealso{ Other ops: -\code{\link[=ops_nms]{ops_nms()}} +\code{\link[=ops_nms]{ops_nms()}}, +\code{\link[=ops_roi_align_rotated]{ops_roi_align_rotated()}} } \concept{ops} diff --git a/man/ops_nms.Rd b/man/ops_nms.Rd index 9ced769..84f6229 100644 --- a/man/ops_nms.Rd +++ b/man/ops_nms.Rd @@ -38,6 +38,7 @@ if (torchvisionlib_is_installed()) { } \seealso{ Other ops: -\code{\link[=ops_ms_deform_attn]{ops_ms_deform_attn()}} +\code{\link[=ops_ms_deform_attn]{ops_ms_deform_attn()}}, +\code{\link[=ops_roi_align_rotated]{ops_roi_align_rotated()}} } \concept{ops} diff --git a/man/ops_roi_align_rotated.Rd b/man/ops_roi_align_rotated.Rd new file mode 100644 index 0000000..7dd7f1f --- /dev/null +++ b/man/ops_roi_align_rotated.Rd @@ -0,0 +1,86 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/ops.R +\name{ops_roi_align_rotated} +\alias{ops_roi_align_rotated} +\alias{nn_roi_align_rotated} +\title{RoI align pooling for rotated proposals} +\usage{ +ops_roi_align_rotated( + input, + rois, + output_size, + spatial_scale, + sampling_ratio = 0, + aligned = TRUE, + clockwise = FALSE +) + +nn_roi_align_rotated( + output_size, + spatial_scale, + sampling_ratio = 0, + aligned = TRUE, + clockwise = FALSE +) +} +\arguments{ +\item{input}{(\code{Tensor[N, C, H, W]}): input feature map.} + +\item{rois}{(\code{Tensor[K, 6]}): rotated boxes with columns +\code{(batch_index, cx, cy, w, h, angle)}, where \code{batch_index} is a \strong{0-based} +index into the first dimension of \code{input}, \code{(cx, cy)} is the box center, +\code{(w, h)} the box size and \code{angle} the rotation angle in radians +(counterclockwise unless \code{clockwise = TRUE}).} + +\item{output_size}{(int or \code{Tuple[int, int]}): the output size \code{(height, width)} +after pooling.} + +\item{spatial_scale}{(float): scaling factor mapping box coordinates to input +coordinates. For example, if boxes are defined on a 224x224 image and +\code{input} is a 112x112 feature map, set this to 0.5.} + +\item{sampling_ratio}{(int): number of sampling points per output bin. If \verb{<= 0}, +an adaptive number of grid points is used (computed as \code{ceil(roi_size / output_size)}). +Default: 0} + +\item{aligned}{(bool): if \code{TRUE} (default), the aligned implementation is +used (results are shifted by -0.5 before interpolation, matching +detectron2). If \code{FALSE}, the legacy MMDetection implementation is used and +boxes are clamped to a minimum size of 1.} + +\item{clockwise}{(bool): if \code{TRUE}, the rotation angle is interpreted in a +clockwise fashion in image space, otherwise it is counterclockwise. +Default: \code{FALSE}} +} +\value{ +\code{Tensor[K, C, output_size[1], output_size[2]]}: the pooled features, where +the \code{r}-th element corresponds to the \code{r}-th RoI in \code{rois}. +} +\description{ +Performs RoI align pooling for rotated proposals, as implemented by the MMCV +\code{roi_align_rotated} operator +(see \url{https://mmcv.readthedocs.io/en/latest/deployment/mmcv_ops_definition.html#mmcvroialignrotated}). +Only a CPU implementation is provided. +} +\section{Functions}{ +\itemize{ +\item \code{nn_roi_align_rotated()}: The \code{\link[torch:nn_module]{torch::nn_module()}} wrapper for \code{\link[=ops_roi_align_rotated]{ops_roi_align_rotated()}}. + +}} +\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) +} + +} +\seealso{ +Other ops: +\code{\link[=ops_ms_deform_attn]{ops_ms_deform_attn()}}, +\code{\link[=ops_nms]{ops_nms()}} +} +\concept{ops} diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index 2deed52..f64768c 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -40,6 +40,24 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// rcpp_vision_ops_roi_align_rotated +torch::Tensor rcpp_vision_ops_roi_align_rotated(torch::Tensor input, torch::Tensor rois, std::int64_t pooled_height, std::int64_t pooled_width, double spatial_scale, std::int64_t sampling_ratio, bool aligned, bool clockwise); +RcppExport SEXP _torchvisionlib_rcpp_vision_ops_roi_align_rotated(SEXP inputSEXP, SEXP roisSEXP, SEXP pooled_heightSEXP, SEXP pooled_widthSEXP, SEXP spatial_scaleSEXP, SEXP sampling_ratioSEXP, SEXP alignedSEXP, SEXP clockwiseSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< torch::Tensor >::type input(inputSEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type rois(roisSEXP); + Rcpp::traits::input_parameter< std::int64_t >::type pooled_height(pooled_heightSEXP); + Rcpp::traits::input_parameter< std::int64_t >::type pooled_width(pooled_widthSEXP); + Rcpp::traits::input_parameter< double >::type spatial_scale(spatial_scaleSEXP); + Rcpp::traits::input_parameter< std::int64_t >::type sampling_ratio(sampling_ratioSEXP); + Rcpp::traits::input_parameter< bool >::type aligned(alignedSEXP); + Rcpp::traits::input_parameter< bool >::type clockwise(clockwiseSEXP); + rcpp_result_gen = Rcpp::wrap(rcpp_vision_ops_roi_align_rotated(input, rois, pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise)); + return rcpp_result_gen; +END_RCPP +} // rcpp_vision_ops_deform_conv2d torch::Tensor rcpp_vision_ops_deform_conv2d(torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor 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); RcppExport SEXP _torchvisionlib_rcpp_vision_ops_deform_conv2d(SEXP inputSEXP, SEXP weightSEXP, SEXP offsetSEXP, SEXP maskSEXP, SEXP biasSEXP, SEXP stride_hSEXP, SEXP stride_wSEXP, SEXP pad_hSEXP, SEXP pad_wSEXP, SEXP dilation_hSEXP, SEXP dilation_wSEXP, SEXP groupsSEXP, SEXP offset_groupsSEXP, SEXP use_maskSEXP) { @@ -185,6 +203,7 @@ END_RCPP static const R_CallMethodDef CallEntries[] = { {"_torchvisionlib_rcpp_vision_ops_nms", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_nms, 3}, {"_torchvisionlib_rcpp_vision_ops_ms_deform_attn", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ms_deform_attn, 6}, + {"_torchvisionlib_rcpp_vision_ops_roi_align_rotated", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_roi_align_rotated, 8}, {"_torchvisionlib_rcpp_vision_ops_deform_conv2d", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_deform_conv2d, 14}, {"_torchvisionlib_rcpp_vision_ops_ps_roi_align", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ps_roi_align, 6}, {"_torchvisionlib_rcpp_vision_ops_ps_roi_pool", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_ps_roi_pool, 5}, diff --git a/src/exports.cpp b/src/exports.cpp index f508f90..4ac7f52 100644 --- a/src/exports.cpp +++ b/src/exports.cpp @@ -12,6 +12,10 @@ torch::Tensor rcpp_vision_ops_ms_deform_attn (torch::Tensor value, torch::Tensor return vision_ops_ms_deform_attn(value.get(), spatial_shapes.get(), level_start_index.get(), sampling_loc.get(), attn_weight.get(), im2col_step); } // [[Rcpp::export]] +torch::Tensor rcpp_vision_ops_roi_align_rotated (torch::Tensor input, torch::Tensor rois, std::int64_t pooled_height, std::int64_t pooled_width, double spatial_scale, std::int64_t sampling_ratio, bool aligned, bool clockwise) { + return vision_ops_roi_align_rotated(input.get(), rois.get(), pooled_height, pooled_width, spatial_scale, sampling_ratio, aligned, clockwise); +} +// [[Rcpp::export]] torch::Tensor rcpp_vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor 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) { return vision_ops_deform_conv2d(input.get(), weight.get(), offset.get(), mask.get(), bias.get(), stride_h, stride_w, pad_h, pad_w, dilation_h, dilation_w, groups, offset_groups, use_mask); } diff --git a/src/exports.h b/src/exports.h index 4444e58..1b05ffe 100644 --- a/src/exports.h +++ b/src/exports.h @@ -5,6 +5,7 @@ torch::Tensor rcpp_vision_ops_nms (torch::Tensor dets, torch::Tensor scores, double iou_threshold); torch::Tensor rcpp_vision_ops_ms_deform_attn (torch::Tensor value, torch::Tensor spatial_shapes, torch::Tensor level_start_index, torch::Tensor sampling_loc, torch::Tensor attn_weight, std::int64_t im2col_step); +torch::Tensor rcpp_vision_ops_roi_align_rotated (torch::Tensor input, torch::Tensor rois, std::int64_t pooled_height, std::int64_t pooled_width, double spatial_scale, std::int64_t sampling_ratio, bool aligned, bool clockwise); torch::Tensor rcpp_vision_ops_deform_conv2d (torch::Tensor input, torch::Tensor weight, torch::Tensor offset, torch::Tensor mask, torch::Tensor 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::tensor_pair rcpp_vision_ops_ps_roi_align (torch::Tensor input, torch::Tensor rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width, int64_t sampling_ratio); torchvisionlib::tensor_pair rcpp_vision_ops_ps_roi_pool (torch::Tensor input, torch::Tensor rois, double spatial_scale, int64_t pooled_height, int64_t pooled_width); diff --git a/tests/run-roi-align-rotated.R b/tests/run-roi-align-rotated.R new file mode 100644 index 0000000..f8b2993 --- /dev/null +++ b/tests/run-roi-align-rotated.R @@ -0,0 +1,197 @@ +# Standalone test for ops_roi_align_rotated (mmcv roi_align_rotated, CPU). +# +# Run after installing torchvisionlib on Windows: +# Rscript tests/run-roi-align-rotated.R +# +# Prints PASS/FAIL for each check and exits with status 0/1. + +suppressMessages({ library(torch); library(torchvisionlib) }) + +pass <- 0L +fail <- 0L + +check <- function(label, cond) { + if (isTRUE(cond)) { + pass <<- pass + 1L + cat(sprintf("PASS %s\n", label)) + } else { + fail <<- fail + 1L + cat(sprintf("FAIL %s\n", label)) + } +} + +# --------------------------------------------------------------------------- +# Pure-R reference of the mmcv CPU kernel +# --------------------------------------------------------------------------- +bilinear_sample <- function(feat, ys, xs) { + H <- dim(feat)[1]; W <- dim(feat)[2] + n <- length(ys) + out <- numeric(n) + for (i in seq_len(n)) { + y <- ys[i]; x <- xs[i] + if (y < -1 || y > H || x < -1 || x > W) { out[i] <- 0; next } + y <- max(y, 0); x <- max(x, 0) + y_low <- floor(y); x_low <- floor(x) + if (y_low >= H - 1) { y_high <- y_low; y_low <- H - 1; y <- H - 1 } else { y_high <- y_low + 1 } + if (x_low >= W - 1) { x_high <- x_low; x_low <- W - 1; x <- W - 1 } else { x_high <- x_low + 1 } + ly <- y - y_low; lx <- x - x_low + hy <- 1 - ly; hx <- 1 - lx + out[i] <- hy * hx * feat[y_low + 1, x_low + 1] + + hy * lx * feat[y_low + 1, x_high + 1] + + ly * hx * feat[y_high + 1, x_low + 1] + + ly * lx * feat[y_high + 1, x_high + 1] + } + out +} + +roi_align_rotated_reference <- function(input, rois, output_size, spatial_scale, + sampling_ratio = 0, aligned = TRUE, + clockwise = FALSE) { + N <- dim(input)[1]; C <- dim(input)[2]; H <- dim(input)[3]; W <- dim(input)[4] + out_h <- output_size[1]; out_w <- output_size[2] + K <- dim(rois)[1] + + input_r <- as.array(input) + rois_m <- matrix(as.numeric(rois), ncol = 6) + out <- array(0, dim = c(K, C, out_h, out_w)) + + for (n in seq_len(K)) { + roi <- rois_m[n, ] + roi_batch_ind <- as.integer(roi[1]) + 1L + offset <- if (aligned) 0.5 else 0.0 + roi_center_w <- roi[2] * spatial_scale - offset + roi_center_h <- roi[3] * spatial_scale - offset + roi_width <- roi[4] * spatial_scale + roi_height <- roi[5] * spatial_scale + theta <- roi[6] + if (clockwise) theta <- -theta + cos_theta <- cos(theta) + sin_theta <- sin(theta) + if (!aligned) { roi_width <- max(roi_width, 1); roi_height <- max(roi_height, 1) } + bin_size_h <- roi_height / out_h + bin_size_w <- roi_width / out_w + grid_h <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_height / out_h) + grid_w <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_width / out_w) + count <- max(grid_h * grid_w, 1) + roi_start_h <- -roi_height / 2 + roi_start_w <- -roi_width / 2 + + ys <- xs <- numeric(0) + for (ph in seq_len(out_h)) { + for (pw in seq_len(out_w)) { + for (iy in seq_len(grid_h)) { + yy <- roi_start_h + (ph - 1) * bin_size_h + (iy - 0.5) * bin_size_h / grid_h + for (ix in seq_len(grid_w)) { + xx <- roi_start_w + (pw - 1) * bin_size_w + (ix - 0.5) * bin_size_w / grid_w + y <- yy * cos_theta - xx * sin_theta + roi_center_h + x <- yy * sin_theta + xx * cos_theta + roi_center_w + ys <- c(ys, y) + xs <- c(xs, x) + } + } + } + } + for (c in seq_len(C)) { + vals <- bilinear_sample(input_r[roi_batch_ind, c, , ], ys, xs) + means <- colMeans(matrix(vals, nrow = grid_h * grid_w)) + out[n, c, , ] <- matrix(means, nrow = out_h, byrow = TRUE) + } + } + out +} + +# --------------------------------------------------------------------------- +# 1. Forward matches the reference +# --------------------------------------------------------------------------- +torch::torch_manual_seed(42) +input <- torch_randn(2, 3, 8, 8) +rois <- torch_tensor( + matrix(c( + 0, 3.5, 3.5, 5, 5, 0.5, # batch 0, ~centered box + 1, 4.2, 4.8, 6, 4, -0.3, # batch 1, wider than tall + 0, 2.0, 2.0, 3, 3, pi / 4 # batch 0, small rotated box + ), ncol = 6, byrow = TRUE), + dtype = torch_float32() +) + +out <- ops_roi_align_rotated(input, rois, c(3, 3), spatial_scale = 1, + sampling_ratio = 2, aligned = TRUE) +ref <- roi_align_rotated_reference(input, rois, c(3, 3), 1, 2, TRUE) +check("forward aligned=TRUE shape", + identical(dim(out), c(3L, 3L, 3L, 3L))) +check("forward aligned=TRUE vs reference", + torch_allclose(out, torch_tensor(ref), atol = 1e-5, rtol = 1e-5)) + +out2 <- ops_roi_align_rotated(input, rois, c(2, 4), spatial_scale = 0.5, + sampling_ratio = 0, aligned = FALSE, clockwise = TRUE) +ref2 <- roi_align_rotated_reference(input, rois, c(2, 4), 0.5, 0, FALSE, TRUE) +check("forward aligned=FALSE/clockwise shape", + identical(dim(out2), c(3L, 3L, 2L, 4L))) +check("forward aligned=FALSE/clockwise vs reference", + torch_allclose(out2, torch_tensor(ref2), atol = 1e-5, rtol = 1e-5)) + +# --------------------------------------------------------------------------- +# 2. Gradients (autograd) match finite differences +# --------------------------------------------------------------------------- +torch::torch_manual_seed(1) +input_g <- torch_randn(2, 2, 6, 6, requires_grad = TRUE) +rois_g <- torch_tensor( + matrix(c(0, 3, 3, 5, 5, 0.4, 1, 3, 4, 4, 4, -0.2), ncol = 6, byrow = TRUE), + dtype = torch_float32() +) +output_size <- c(3, 3) + +ops_roi_align_rotated(input_g, rois_g, output_size, spatial_scale = 1, + sampling_ratio = 2)$sum()$backward() +check("backward produced grad", !is.null(input_g$grad)) + +eps <- 1e-3 +num_grad <- torch_empty_like(input_g) +input_flat <- input_g$detach()$flatten() +for (i in seq_len(prod(dim(input_g)))) { + x_p <- input_flat$clone(); x_m <- input_flat$clone() + x_p[i] <- as.numeric(x_p[i]) + eps + x_m[i] <- as.numeric(x_m[i]) - eps + f_p <- ops_roi_align_rotated(x_p$view(dim(input_g)), rois_g, output_size, 1, + sampling_ratio = 2)$sum() + f_m <- ops_roi_align_rotated(x_m$view(dim(input_g)), rois_g, output_size, 1, + sampling_ratio = 2)$sum() + num_grad$flatten()[i] <- as.numeric(f_p - f_m) / (2 * eps) +} +check("backward matches finite differences", + torch_allclose(input_g$grad, num_grad, atol = 1e-3, rtol = 1e-3)) + +# --------------------------------------------------------------------------- +# 3. nn module +# --------------------------------------------------------------------------- +mod <- nn_roi_align_rotated(output_size = c(4, 4), spatial_scale = 1, sampling_ratio = 1) +out_mod <- mod(torch_randn(1, 2, 10, 10), + torch_tensor(matrix(c(0, 5, 5, 4, 4, 0.3), ncol = 6), dtype = torch_float32())) +check("nn_roi_align_rotated output shape", identical(dim(out_mod), c(1L, 2L, 4L, 4L))) + +# --------------------------------------------------------------------------- +# 4. Input validation +# --------------------------------------------------------------------------- +ok_err <- function(expr, pattern) { + res <- tryCatch({ force(expr); "no error" }, error = function(e) conditionMessage(e)) + grepl(pattern, res, fixed = TRUE) +} +check("error: rois must have 6 columns", + ok_err( + ops_roi_align_rotated(input, torch_tensor(matrix(c(0, 3, 3, 5, 5), ncol = 5), dtype = torch_float32()), c(3, 3), 1), + "rois should have 6 columns")) +check("error: negative roi size when aligned", + ok_err( + ops_roi_align_rotated(input, torch_tensor(matrix(c(0, 3, 3, -2, 5, 0.4), ncol = 6), dtype = torch_float32()), c(3, 3), 1), + "do not have non-negative size")) +check("error: out-of-range batch index", + ok_err( + ops_roi_align_rotated(input, torch_tensor(matrix(c(5, 3, 3, 5, 5, 0.4), ncol = 6), dtype = torch_float32()), c(3, 3), 1), + "rois index should be in [0, batch_size)")) +check("error: non-positive output size", + ok_err(ops_roi_align_rotated(input, rois, c(0, 3), 1), + "pooled_height and pooled_width should be positive")) + +# --------------------------------------------------------------------------- +cat(sprintf("\n%d passed, %d failed\n", pass, fail)) +if (fail > 0) quit(status = 1) diff --git a/tests/testthat/test-ops-roi-align-rotated.R b/tests/testthat/test-ops-roi-align-rotated.R new file mode 100644 index 0000000..b679d9b --- /dev/null +++ b/tests/testthat/test-ops-roi-align-rotated.R @@ -0,0 +1,205 @@ +library(torch) + +# Pure-R reference implementation of mmcv's `roi_align_rotated` CPU kernel +# (https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/pytorch/cpu/roi_align_rotated_cpu.cpp). +# Used to validate the C++ implementation independently. + +bilinear_sample <- function(feat, ys, xs) { + H <- dim(feat)[1]; W <- dim(feat)[2] + n <- length(ys) + out <- numeric(n) + for (i in seq_len(n)) { + y <- ys[i]; x <- xs[i] + if (y < -1 || y > H || x < -1 || x > W) { + out[i] <- 0 + next + } + y <- max(y, 0); x <- max(x, 0) + y_low <- floor(y); x_low <- floor(x) + if (y_low >= H - 1) { y_high <- y_low; y_low <- H - 1; y <- H - 1 } else { y_high <- y_low + 1 } + if (x_low >= W - 1) { x_high <- x_low; x_low <- W - 1; x <- W - 1 } else { x_high <- x_low + 1 } + ly <- y - y_low; lx <- x - x_low + hy <- 1 - ly; hx <- 1 - lx + out[i] <- hy * hx * feat[y_low + 1, x_low + 1] + + hy * lx * feat[y_low + 1, x_high + 1] + + ly * hx * feat[y_high + 1, x_low + 1] + + ly * lx * feat[y_high + 1, x_high + 1] + } + out +} + +roi_align_rotated_reference <- function(input, rois, output_size, spatial_scale, + sampling_ratio = 0, aligned = TRUE, + clockwise = FALSE) { + stopifnot(length(dim(input)) == 4) + N <- dim(input)[1]; C <- dim(input)[2]; H <- dim(input)[3]; W <- dim(input)[4] + out_h <- output_size[1]; out_w <- output_size[2] + K <- dim(rois)[1] + + input_r <- as.array(input) + rois_m <- matrix(as.numeric(rois), ncol = 6) + out <- array(0, dim = c(K, C, out_h, out_w)) + + for (n in seq_len(K)) { + roi <- rois_m[n, ] + roi_batch_ind <- as.integer(roi[1]) + 1L # 0-based in C++, 1-based in R + offset <- if (aligned) 0.5 else 0.0 + roi_center_w <- roi[2] * spatial_scale - offset + roi_center_h <- roi[3] * spatial_scale - offset + roi_width <- roi[4] * spatial_scale + roi_height <- roi[5] * spatial_scale + theta <- roi[6] + if (clockwise) theta <- -theta + cos_theta <- cos(theta) + sin_theta <- sin(theta) + if (!aligned) { + roi_width <- max(roi_width, 1) + roi_height <- max(roi_height, 1) + } + bin_size_h <- roi_height / out_h + bin_size_w <- roi_width / out_w + grid_h <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_height / out_h) + grid_w <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_width / out_w) + count <- max(grid_h * grid_w, 1) + roi_start_h <- -roi_height / 2 + roi_start_w <- -roi_width / 2 + + ys <- xs <- numeric(0) + for (ph in seq_len(out_h)) { + for (pw in seq_len(out_w)) { + for (iy in seq_len(grid_h)) { + yy <- roi_start_h + (ph - 1) * bin_size_h + (iy - 0.5) * bin_size_h / grid_h + for (ix in seq_len(grid_w)) { + xx <- roi_start_w + (pw - 1) * bin_size_w + (ix - 0.5) * bin_size_w / grid_w + y <- yy * cos_theta - xx * sin_theta + roi_center_h + x <- yy * sin_theta + xx * cos_theta + roi_center_w + ys <- c(ys, y) + xs <- c(xs, x) + } + } + } + } + + for (c in seq_len(C)) { + vals <- bilinear_sample(input_r[roi_batch_ind, c, , ], ys, xs) + means <- colMeans(matrix(vals, nrow = grid_h * grid_w)) + out[n, c, , ] <- matrix(means, nrow = out_h, byrow = TRUE) + } + } + out +} + +make_rois <- function() { + torch_tensor( + matrix(c( + 0, 3.5, 3.5, 5, 5, 0.5, # batch 0, ~centered box + 1, 4.2, 4.8, 6, 4, -0.3, # batch 1, wider than tall + 0, 2.0, 2.0, 3, 3, pi / 4 # batch 0, small rotated box + ), ncol = 6, byrow = TRUE), + dtype = torch_float32() + ) +} + +make_input <- function() { + torch::torch_manual_seed(42) + torch_randn(2, 3, 8, 8) +} + +test_that("roi_align_rotated matches the mmcv reference (forward)", { + input <- make_input() + rois <- make_rois() + + out <- ops_roi_align_rotated(input, rois, c(3, 3), spatial_scale = 1, + sampling_ratio = 2, aligned = TRUE) + ref <- roi_align_rotated_reference(input, rois, c(3, 3), 1, 2, TRUE) + expect_equal(dim(out), c(3, 3, 3, 3)) + expect_true(torch_allclose(out, torch_tensor(ref), atol = 1e-5, rtol = 1e-5)) +}) + +test_that("roi_align_rotated matches with aligned=FALSE and sampling_ratio=0", { + input <- make_input() + rois <- make_rois() + + out <- ops_roi_align_rotated(input, rois, c(2, 4), spatial_scale = 0.5, + sampling_ratio = 0, aligned = FALSE, + clockwise = TRUE) + ref <- roi_align_rotated_reference(input, rois, c(2, 4), 0.5, 0, FALSE, TRUE) + expect_equal(dim(out), c(3, 3, 2, 4)) + expect_true(torch_allclose(out, torch_tensor(ref), atol = 1e-5, rtol = 1e-5)) +}) + +test_that("roi_align_rotated is differentiable and matches finite differences", { + input <- torch::torch_randn(2, 2, 6, 6, requires_grad = TRUE) + rois <- torch_tensor( + matrix(c(0, 3, 3, 5, 5, 0.4, 1, 3, 4, 4, 4, -0.2), ncol = 6, byrow = TRUE), + dtype = torch_float32() + ) + output_size <- c(3, 3) + + out <- ops_roi_align_rotated(input, rois, output_size, spatial_scale = 1, + sampling_ratio = 2) + out$sum()$backward() + expect_true(!is.null(input$grad)) + + eps <- 1e-3 + num_grad <- torch_empty_like(input) + input_flat <- input$detach()$flatten() + for (i in seq_len(numel <- prod(dim(input)))) { + x_p <- input_flat$clone() + x_m <- input_flat$clone() + x_p[i] <- as.numeric(x_p[i]) + eps + x_m[i] <- as.numeric(x_m[i]) - eps + f_p <- ops_roi_align_rotated(x_p$view(dim(input)), rois, output_size, 1, + sampling_ratio = 2)$sum() + f_m <- ops_roi_align_rotated(x_m$view(dim(input)), rois, output_size, 1, + sampling_ratio = 2)$sum() + num_grad$flatten()[i] <- as.numeric(f_p - f_m) / (2 * eps) + } + + expect_true(torch_allclose(input$grad, num_grad, atol = 1e-3, rtol = 1e-3)) +}) + +test_that("nn_roi_align_rotated module works", { + input <- torch_randn(1, 2, 10, 10) + rois <- torch_tensor(matrix(c(0, 5, 5, 4, 4, 0.3), ncol = 6), + dtype = torch_float32()) + mod <- nn_roi_align_rotated(output_size = c(4, 4), spatial_scale = 1, + sampling_ratio = 1) + out <- mod(input, rois) + expect_equal(dim(out), c(1, 2, 4, 4)) +}) + +test_that("roi_align_rotated validates its inputs", { + input <- make_input() + rois <- make_rois() + + # wrong number of rois columns + bad_rois <- torch_tensor(matrix(c(0, 3, 3, 5, 5), ncol = 5), + dtype = torch_float32()) + expect_error( + ops_roi_align_rotated(input, bad_rois, c(3, 3), 1), + regexp = "rois should have 6 columns" + ) + + # negative box size with aligned = TRUE + neg_rois <- torch_tensor(matrix(c(0, 3, 3, -2, 5, 0.4), ncol = 6), + dtype = torch_float32()) + expect_error( + ops_roi_align_rotated(input, neg_rois, c(3, 3), 1), + regexp = "do not have non-negative size" + ) + + # out-of-range batch index + oob_rois <- torch_tensor(matrix(c(5, 3, 3, 5, 5, 0.4), ncol = 6), + dtype = torch_float32()) + expect_error( + ops_roi_align_rotated(input, oob_rois, c(3, 3), 1), + regexp = "rois index should be in \\[0, batch_size\\)" + ) + + # non-positive output size + expect_error( + ops_roi_align_rotated(input, rois, c(0, 3), 1), + regexp = "pooled_height and pooled_width should be positive" + ) +}) From 0167957288e01707ecc6cf824fb2ebe3a3218ad9 Mon Sep 17 00:00:00 2001 From: DerrickUnleashed Date: Mon, 3 Aug 2026 00:11:44 +0530 Subject: [PATCH 2/4] Sync inst/include exports header with csrc (#32) --- inst/include/torchvisionlib/exports.h | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/inst/include/torchvisionlib/exports.h b/inst/include/torchvisionlib/exports.h index 16278a9..b991ac8 100644 --- a/inst/include/torchvisionlib/exports.h +++ b/inst/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(); From 0b8fa3342aae24b638d9705a088aabdbf0b876e6 Mon Sep 17 00:00:00 2001 From: DerrickUnleashed Date: Mon, 3 Aug 2026 00:30:36 +0530 Subject: [PATCH 3/4] Replace and rename test file according to conventions --- .../test-roi-align-rotated.R} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{run-roi-align-rotated.R => testthat/test-roi-align-rotated.R} (100%) diff --git a/tests/run-roi-align-rotated.R b/tests/testthat/test-roi-align-rotated.R similarity index 100% rename from tests/run-roi-align-rotated.R rename to tests/testthat/test-roi-align-rotated.R From ed6a3fe167dda1549cfe03d6dbade604307e4068 Mon Sep 17 00:00:00 2001 From: DerrickUnleashed Date: Mon, 3 Aug 2026 02:04:30 +0530 Subject: [PATCH 4/4] Fix gradient check in roi_align_rotated test and drop duplicate --- tests/testthat/test-ops-roi-align-rotated.R | 3 +- tests/testthat/test-roi-align-rotated.R | 197 -------------------- 2 files changed, 2 insertions(+), 198 deletions(-) delete mode 100644 tests/testthat/test-roi-align-rotated.R diff --git a/tests/testthat/test-ops-roi-align-rotated.R b/tests/testthat/test-ops-roi-align-rotated.R index b679d9b..f671cf2 100644 --- a/tests/testthat/test-ops-roi-align-rotated.R +++ b/tests/testthat/test-ops-roi-align-rotated.R @@ -143,6 +143,7 @@ test_that("roi_align_rotated is differentiable and matches finite differences", eps <- 1e-3 num_grad <- torch_empty_like(input) + num_grad_flat <- num_grad$flatten() input_flat <- input$detach()$flatten() for (i in seq_len(numel <- prod(dim(input)))) { x_p <- input_flat$clone() @@ -153,7 +154,7 @@ test_that("roi_align_rotated is differentiable and matches finite differences", sampling_ratio = 2)$sum() f_m <- ops_roi_align_rotated(x_m$view(dim(input)), rois, output_size, 1, sampling_ratio = 2)$sum() - num_grad$flatten()[i] <- as.numeric(f_p - f_m) / (2 * eps) + num_grad_flat[i] <- as.numeric(f_p - f_m) / (2 * eps) } expect_true(torch_allclose(input$grad, num_grad, atol = 1e-3, rtol = 1e-3)) diff --git a/tests/testthat/test-roi-align-rotated.R b/tests/testthat/test-roi-align-rotated.R deleted file mode 100644 index f8b2993..0000000 --- a/tests/testthat/test-roi-align-rotated.R +++ /dev/null @@ -1,197 +0,0 @@ -# Standalone test for ops_roi_align_rotated (mmcv roi_align_rotated, CPU). -# -# Run after installing torchvisionlib on Windows: -# Rscript tests/run-roi-align-rotated.R -# -# Prints PASS/FAIL for each check and exits with status 0/1. - -suppressMessages({ library(torch); library(torchvisionlib) }) - -pass <- 0L -fail <- 0L - -check <- function(label, cond) { - if (isTRUE(cond)) { - pass <<- pass + 1L - cat(sprintf("PASS %s\n", label)) - } else { - fail <<- fail + 1L - cat(sprintf("FAIL %s\n", label)) - } -} - -# --------------------------------------------------------------------------- -# Pure-R reference of the mmcv CPU kernel -# --------------------------------------------------------------------------- -bilinear_sample <- function(feat, ys, xs) { - H <- dim(feat)[1]; W <- dim(feat)[2] - n <- length(ys) - out <- numeric(n) - for (i in seq_len(n)) { - y <- ys[i]; x <- xs[i] - if (y < -1 || y > H || x < -1 || x > W) { out[i] <- 0; next } - y <- max(y, 0); x <- max(x, 0) - y_low <- floor(y); x_low <- floor(x) - if (y_low >= H - 1) { y_high <- y_low; y_low <- H - 1; y <- H - 1 } else { y_high <- y_low + 1 } - if (x_low >= W - 1) { x_high <- x_low; x_low <- W - 1; x <- W - 1 } else { x_high <- x_low + 1 } - ly <- y - y_low; lx <- x - x_low - hy <- 1 - ly; hx <- 1 - lx - out[i] <- hy * hx * feat[y_low + 1, x_low + 1] + - hy * lx * feat[y_low + 1, x_high + 1] + - ly * hx * feat[y_high + 1, x_low + 1] + - ly * lx * feat[y_high + 1, x_high + 1] - } - out -} - -roi_align_rotated_reference <- function(input, rois, output_size, spatial_scale, - sampling_ratio = 0, aligned = TRUE, - clockwise = FALSE) { - N <- dim(input)[1]; C <- dim(input)[2]; H <- dim(input)[3]; W <- dim(input)[4] - out_h <- output_size[1]; out_w <- output_size[2] - K <- dim(rois)[1] - - input_r <- as.array(input) - rois_m <- matrix(as.numeric(rois), ncol = 6) - out <- array(0, dim = c(K, C, out_h, out_w)) - - for (n in seq_len(K)) { - roi <- rois_m[n, ] - roi_batch_ind <- as.integer(roi[1]) + 1L - offset <- if (aligned) 0.5 else 0.0 - roi_center_w <- roi[2] * spatial_scale - offset - roi_center_h <- roi[3] * spatial_scale - offset - roi_width <- roi[4] * spatial_scale - roi_height <- roi[5] * spatial_scale - theta <- roi[6] - if (clockwise) theta <- -theta - cos_theta <- cos(theta) - sin_theta <- sin(theta) - if (!aligned) { roi_width <- max(roi_width, 1); roi_height <- max(roi_height, 1) } - bin_size_h <- roi_height / out_h - bin_size_w <- roi_width / out_w - grid_h <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_height / out_h) - grid_w <- if (sampling_ratio > 0) sampling_ratio else ceiling(roi_width / out_w) - count <- max(grid_h * grid_w, 1) - roi_start_h <- -roi_height / 2 - roi_start_w <- -roi_width / 2 - - ys <- xs <- numeric(0) - for (ph in seq_len(out_h)) { - for (pw in seq_len(out_w)) { - for (iy in seq_len(grid_h)) { - yy <- roi_start_h + (ph - 1) * bin_size_h + (iy - 0.5) * bin_size_h / grid_h - for (ix in seq_len(grid_w)) { - xx <- roi_start_w + (pw - 1) * bin_size_w + (ix - 0.5) * bin_size_w / grid_w - y <- yy * cos_theta - xx * sin_theta + roi_center_h - x <- yy * sin_theta + xx * cos_theta + roi_center_w - ys <- c(ys, y) - xs <- c(xs, x) - } - } - } - } - for (c in seq_len(C)) { - vals <- bilinear_sample(input_r[roi_batch_ind, c, , ], ys, xs) - means <- colMeans(matrix(vals, nrow = grid_h * grid_w)) - out[n, c, , ] <- matrix(means, nrow = out_h, byrow = TRUE) - } - } - out -} - -# --------------------------------------------------------------------------- -# 1. Forward matches the reference -# --------------------------------------------------------------------------- -torch::torch_manual_seed(42) -input <- torch_randn(2, 3, 8, 8) -rois <- torch_tensor( - matrix(c( - 0, 3.5, 3.5, 5, 5, 0.5, # batch 0, ~centered box - 1, 4.2, 4.8, 6, 4, -0.3, # batch 1, wider than tall - 0, 2.0, 2.0, 3, 3, pi / 4 # batch 0, small rotated box - ), ncol = 6, byrow = TRUE), - dtype = torch_float32() -) - -out <- ops_roi_align_rotated(input, rois, c(3, 3), spatial_scale = 1, - sampling_ratio = 2, aligned = TRUE) -ref <- roi_align_rotated_reference(input, rois, c(3, 3), 1, 2, TRUE) -check("forward aligned=TRUE shape", - identical(dim(out), c(3L, 3L, 3L, 3L))) -check("forward aligned=TRUE vs reference", - torch_allclose(out, torch_tensor(ref), atol = 1e-5, rtol = 1e-5)) - -out2 <- ops_roi_align_rotated(input, rois, c(2, 4), spatial_scale = 0.5, - sampling_ratio = 0, aligned = FALSE, clockwise = TRUE) -ref2 <- roi_align_rotated_reference(input, rois, c(2, 4), 0.5, 0, FALSE, TRUE) -check("forward aligned=FALSE/clockwise shape", - identical(dim(out2), c(3L, 3L, 2L, 4L))) -check("forward aligned=FALSE/clockwise vs reference", - torch_allclose(out2, torch_tensor(ref2), atol = 1e-5, rtol = 1e-5)) - -# --------------------------------------------------------------------------- -# 2. Gradients (autograd) match finite differences -# --------------------------------------------------------------------------- -torch::torch_manual_seed(1) -input_g <- torch_randn(2, 2, 6, 6, requires_grad = TRUE) -rois_g <- torch_tensor( - matrix(c(0, 3, 3, 5, 5, 0.4, 1, 3, 4, 4, 4, -0.2), ncol = 6, byrow = TRUE), - dtype = torch_float32() -) -output_size <- c(3, 3) - -ops_roi_align_rotated(input_g, rois_g, output_size, spatial_scale = 1, - sampling_ratio = 2)$sum()$backward() -check("backward produced grad", !is.null(input_g$grad)) - -eps <- 1e-3 -num_grad <- torch_empty_like(input_g) -input_flat <- input_g$detach()$flatten() -for (i in seq_len(prod(dim(input_g)))) { - x_p <- input_flat$clone(); x_m <- input_flat$clone() - x_p[i] <- as.numeric(x_p[i]) + eps - x_m[i] <- as.numeric(x_m[i]) - eps - f_p <- ops_roi_align_rotated(x_p$view(dim(input_g)), rois_g, output_size, 1, - sampling_ratio = 2)$sum() - f_m <- ops_roi_align_rotated(x_m$view(dim(input_g)), rois_g, output_size, 1, - sampling_ratio = 2)$sum() - num_grad$flatten()[i] <- as.numeric(f_p - f_m) / (2 * eps) -} -check("backward matches finite differences", - torch_allclose(input_g$grad, num_grad, atol = 1e-3, rtol = 1e-3)) - -# --------------------------------------------------------------------------- -# 3. nn module -# --------------------------------------------------------------------------- -mod <- nn_roi_align_rotated(output_size = c(4, 4), spatial_scale = 1, sampling_ratio = 1) -out_mod <- mod(torch_randn(1, 2, 10, 10), - torch_tensor(matrix(c(0, 5, 5, 4, 4, 0.3), ncol = 6), dtype = torch_float32())) -check("nn_roi_align_rotated output shape", identical(dim(out_mod), c(1L, 2L, 4L, 4L))) - -# --------------------------------------------------------------------------- -# 4. Input validation -# --------------------------------------------------------------------------- -ok_err <- function(expr, pattern) { - res <- tryCatch({ force(expr); "no error" }, error = function(e) conditionMessage(e)) - grepl(pattern, res, fixed = TRUE) -} -check("error: rois must have 6 columns", - ok_err( - ops_roi_align_rotated(input, torch_tensor(matrix(c(0, 3, 3, 5, 5), ncol = 5), dtype = torch_float32()), c(3, 3), 1), - "rois should have 6 columns")) -check("error: negative roi size when aligned", - ok_err( - ops_roi_align_rotated(input, torch_tensor(matrix(c(0, 3, 3, -2, 5, 0.4), ncol = 6), dtype = torch_float32()), c(3, 3), 1), - "do not have non-negative size")) -check("error: out-of-range batch index", - ok_err( - ops_roi_align_rotated(input, torch_tensor(matrix(c(5, 3, 3, 5, 5, 0.4), ncol = 6), dtype = torch_float32()), c(3, 3), 1), - "rois index should be in [0, batch_size)")) -check("error: non-positive output size", - ok_err(ops_roi_align_rotated(input, rois, c(0, 3), 1), - "pooled_height and pooled_width should be positive")) - -# --------------------------------------------------------------------------- -cat(sprintf("\n%d passed, %d failed\n", pass, fail)) -if (fail > 0) quit(status = 1)