Add box_iou_rotated

LICENSE.md

@@ -19,3 +19,11 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
+
+---
+
+The rotated box IoU kernel under `csrc/ops/box_iou_rotated/` is adapted from
+Detectron2 (https://github.com/facebookresearch/detectron2), Copyright (c)
+Facebook, Inc. and its affiliates, licensed under the Apache License, Version
+2.0, and from Meta's torchvision (BSD-style license). The original license
+headers are retained in those files.

NAMESPACE

@@ -2,6 +2,7 @@
 
 export(install_torchvisionlib)
 export(nn_ps_roi_align)
+export(ops_box_iou_rotated)
 export(ops_deform_conv2d)
 export(ops_nms)
 export(ops_ps_roi_align)

NEWS.md

@@ -1,5 +1,9 @@
 # torchvisionlib (development version)
 
+- Added `ops_box_iou_rotated()`, a CPU implementation of intersection-over-union
+  between rotated boxes, supporting the `cxcywhr`, `xywhr` and `xyxyxyxy`
+  formats. Adapted from Detectron2 (Apache-2.0). (#31)
+
 # torchvisionlib 0.8.0
 
 - Updates to support LibTorch v2.8

R/RcppExports.R

@@ -5,6 +5,10 @@ rcpp_vision_ops_nms <- function(dets, scores, iou_threshold) {
     .Call('_torchvisionlib_rcpp_vision_ops_nms', PACKAGE = 'torchvisionlib', dets, scores, iou_threshold)
 }
 
+rcpp_vision_ops_box_iou_rotated <- function(boxes1, boxes2) {
+    .Call('_torchvisionlib_rcpp_vision_ops_box_iou_rotated', PACKAGE = 'torchvisionlib', boxes1, boxes2)
+}
+
 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

@@ -33,6 +33,37 @@ ops_nms <- function(boxes, scores, iou_threshold) {
 }
 
 
+#' Intersection-over-union between rotated boxes
+#'
+#' Computes the pairwise intersection-over-union (IoU) between two sets of
+#' rotated bounding boxes.
+#'
+#' @param boxes1 `Tensor[N, K]` first set of rotated boxes.
+#' @param boxes2 `Tensor[M, K]` second set of rotated boxes.
+#' @param fmt format of the input boxes. One of:
+#'   * `"cxcywhr"` (`K = 5`): center `(cx, cy)`, width, height and rotation
+#'     angle `r` in degrees (counter-clockwise positive).
+#'   * `"xywhr"` (`K = 5`): top-left corner `(x1, y1)`, width, height and angle.
+#'   * `"xyxyxyxy"` (`K = 8`): the four corners
+#'     `(x1, y1, x2, y2, x3, y3, x4, y4)`.
+#'
+#' @returns
+#' `Tensor[N, M]` float32 matrix of pairwise IoU values.
+#'
+#' @examples
+#' if (torchvisionlib_is_installed()) {
+#'   boxes <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 45), nrow = 1))
+#'   ops_box_iou_rotated(boxes, boxes)
+#' }
+#' @family ops
+#' @export
+ops_box_iou_rotated <- function(boxes1, boxes2, fmt = "cxcywhr") {
+  boxes1 <- .rotated_boxes_to_cxcywhr(boxes1, fmt)
+  boxes2 <- .rotated_boxes_to_cxcywhr(boxes2, fmt)
+  rcpp_vision_ops_box_iou_rotated(boxes1, boxes2)
+}
+
+
 #' Performs Deformable Convolution v2,
 #'
 #' Ddescribed in [Deformable ConvNets v2: More Deformable, Better Results](https://arxiv.org/abs/1811.11168)

R/utils.R

@@ -21,6 +21,50 @@ runtime_error <- function(...) {
   rlang::abort(..., class = "runtime_error")
 }
 
+# Convert a set of rotated boxes to `cxcywhr`, the format expected by the
+# `box_iou_rotated` C++ op. Conversions mirror torchvision's `box_convert`.
+.rotated_boxes_to_cxcywhr <- function(boxes, fmt) {
+  switch(
+    fmt,
+    cxcywhr = boxes,
+    xywhr = .box_xywhr_to_cxcywhr(boxes),
+    xyxyxyxy = .box_xyxyxyxy_to_cxcywhr(boxes),
+    runtime_error(sprintf(
+      "Unsupported format '%s'. Supported rotated formats: cxcywhr, xywhr, xyxyxyxy.",
+      fmt
+    ))
+  )
+}
+
+.box_xywhr_to_cxcywhr <- function(boxes) {
+  b <- torch::torch_unbind(boxes, dim = -1)
+  x1 <- b[[1]]
+  y1 <- b[[2]]
+  w <- b[[3]]
+  h <- b[[4]]
+  r <- b[[5]]
+  r_rad <- r * pi / 180
+  cos <- torch::torch_cos(r_rad)
+  sin <- torch::torch_sin(r_rad)
+  cx <- x1 + w / 2 * cos + h / 2 * sin
+  cy <- y1 - w / 2 * sin + h / 2 * cos
+  torch::torch_stack(list(cx, cy, w, h, r), dim = -1)
+}
+
+.box_xyxyxyxy_to_cxcywhr <- function(boxes) {
+  b <- torch::torch_unbind(boxes, dim = -1)
+  x1 <- b[[1]]
+  y1 <- b[[2]]
+  x2 <- b[[3]]
+  y2 <- b[[4]]
+  x3 <- b[[5]]
+  y3 <- b[[6]]
+  r <- torch::torch_atan2(y1 - y2, x2 - x1) * 180 / pi
+  w <- ((x2 - x1) * (x2 - x1) + (y1 - y2) * (y1 - y2))$sqrt()
+  h <- ((x3 - x2) * (x3 - x2) + (y3 - y2) * (y3 - y2))$sqrt()
+  .box_xywhr_to_cxcywhr(torch::torch_stack(list(x1, y1, w, h, r), dim = -1))
+}
+
 # Efficient version of torch.cat that avoids a copy if there is only a single element in a list
 .cat <- function(tensors, dim = 1) {
   if (length(tensors) == 1)

csrc/CMakeLists.txt

@@ -136,12 +136,15 @@ if (WIN32)
   set_target_properties(TorchVision PROPERTIES IMPORTED_IMPLIB ${TorchVision_DESTDIR}/lib/torchvision.lib)
 endif()
 
-set(TORCHVISION_SRC src/torchvisionlib.cpp src/ops.cpp src/exports.cpp src/torchvisionlib_types.cpp)
+set(TORCHVISION_SRC src/torchvisionlib.cpp src/ops.cpp src/exports.cpp src/torchvisionlib_types.cpp
+    ops/box_iou_rotated/box_iou_rotated.cpp
+    ops/box_iou_rotated/cpu/box_iou_rotated_kernel.cpp)
 
 add_library(torchvisionlib SHARED ${TORCHVISION_SRC})
 add_library(torchvisionlib::library ALIAS torchvisionlib)
 
 target_include_directories(torchvisionlib PUBLIC
+    ${PROJECT_SOURCE_DIR}
     ${PROJECT_SOURCE_DIR}/include
     ${TORCH_HOME}/include
     ${TORCHVISION_INCLUDE_DIR}

csrc/include/torchvisionlib/exports.h

@@ -28,6 +28,7 @@ TORCHVISIONLIB_API void* torchvisionlib_last_error ();
 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_box_iou_rotated (void* boxes1, void* boxes2);
 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);
@@ -45,6 +46,11 @@ inline void* vision_ops_nms (void* dets, void* scores, double iou_threshold) {
   host_exception_handler();
   return ret;
 }
+inline void* vision_ops_box_iou_rotated (void* boxes1, void* boxes2) {
+  auto ret =  _vision_ops_box_iou_rotated(boxes1, boxes2);
+  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/box_iou_rotated/box_iou_rotated.cpp

@@ -0,0 +1,28 @@
+#include "box_iou_rotated.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/library.h>
+#include <torch/types.h>
+
+namespace vision {
+namespace ops {
+
+at::Tensor box_iou_rotated(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("torchvision::box_iou_rotated", "")
+                       .typed<decltype(box_iou_rotated)>();
+  return op.call(boxes1, boxes2);
+}
+
+// Vendored because the pinned TorchVision (v0.20.1) has no box_iou_rotated.
+// Drop this directory if TorchVision is bumped to a release that ships the op,
+// otherwise the schema below is registered twice and loading fails.
+TORCH_LIBRARY_FRAGMENT(torchvision, m) {
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "torchvision::box_iou_rotated(Tensor boxes1, Tensor boxes2) -> Tensor"));
+}
+
+} // namespace ops
+} // namespace vision

csrc/ops/box_iou_rotated/box_iou_rotated.h

@@ -0,0 +1,13 @@
+#pragma once
+
+#include <ATen/ATen.h>
+
+namespace vision {
+namespace ops {
+
+at::Tensor box_iou_rotated(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2);
+
+} // namespace ops
+} // namespace vision