diff --git a/LICENSE.md b/LICENSE.md index 9acb9f0..6175781 100644 --- a/LICENSE.md +++ b/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. diff --git a/NAMESPACE b/NAMESPACE index 5a141c9..8331d9d 100644 --- a/NAMESPACE +++ b/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) diff --git a/NEWS.md b/NEWS.md index bdee967..8cb9734 100644 --- a/NEWS.md +++ b/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 diff --git a/R/RcppExports.R b/R/RcppExports.R index 8b5dd00..ce6c09a 100644 --- a/R/RcppExports.R +++ b/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) } diff --git a/R/ops.R b/R/ops.R index c83c475..7d35d79 100644 --- a/R/ops.R +++ b/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) diff --git a/R/utils.R b/R/utils.R index 5c754ad..a053396 100644 --- a/R/utils.R +++ b/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) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index a6209ee..974e0d0 100644 --- a/csrc/CMakeLists.txt +++ b/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} diff --git a/csrc/include/torchvisionlib/exports.h b/csrc/include/torchvisionlib/exports.h index a8de27c..47f42da 100644 --- a/csrc/include/torchvisionlib/exports.h +++ b/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(); diff --git a/csrc/ops/box_iou_rotated/box_iou_rotated.cpp b/csrc/ops/box_iou_rotated/box_iou_rotated.cpp new file mode 100644 index 0000000..82a806f --- /dev/null +++ b/csrc/ops/box_iou_rotated/box_iou_rotated.cpp @@ -0,0 +1,28 @@ +#include "box_iou_rotated.h" + +#include +#include +#include + +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(); + return op.call(boxes1, boxes2); +} + +// Vendored because the pinned TorchVision (v0.23.0) 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 diff --git a/csrc/ops/box_iou_rotated/box_iou_rotated.h b/csrc/ops/box_iou_rotated/box_iou_rotated.h new file mode 100644 index 0000000..f09f2e5 --- /dev/null +++ b/csrc/ops/box_iou_rotated/box_iou_rotated.h @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace vision { +namespace ops { + +at::Tensor box_iou_rotated( + const at::Tensor& boxes1, + const at::Tensor& boxes2); + +} // namespace ops +} // namespace vision diff --git a/csrc/ops/box_iou_rotated/cpu/box_iou_rotated_kernel.cpp b/csrc/ops/box_iou_rotated/cpu/box_iou_rotated_kernel.cpp new file mode 100644 index 0000000..dcc034e --- /dev/null +++ b/csrc/ops/box_iou_rotated/cpu/box_iou_rotated_kernel.cpp @@ -0,0 +1,412 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. +// +// This file contains code adapted from Detectron2's box_iou_rotated +// implementation, which is licensed under the Apache License, Version 2.0. +// Original source: https://github.com/facebookresearch/detectron2 +// License: https://github.com/facebookresearch/detectron2/blob/main/LICENSE + +#include +#include + +namespace vision { +namespace ops { + +namespace { + +template +struct RotatedBox { + T x_ctr, y_ctr, w, h, a; +}; + +template +struct Point { + T x, y; + Point(const T& px = 0, const T& py = 0) : x(px), y(py) {} + Point operator+(const Point& p) const { + return Point(x + p.x, y + p.y); + } + Point& operator+=(const Point& p) { + x += p.x; + y += p.y; + return *this; + } + Point operator-(const Point& p) const { + return Point(x - p.x, y - p.y); + } + Point operator*(const T& coeff) const { + return Point(x * coeff, y * coeff); + } +}; + +template +inline T dot_2d(const Point& A, const Point& B) { + return A.x * B.x + A.y * B.y; +} + +template +inline T cross_2d(const Point& A, const Point& B) { + return A.x * B.y - B.x * A.y; +} + +template +inline void get_rotated_vertices(const RotatedBox& box, Point (&pts)[4]) { + // M_PI / 180. == 0.01745329251 + double theta = box.a * 0.01745329251; + T cosTheta2 = (T)cos(theta) * 0.5f; + T sinTheta2 = (T)sin(theta) * 0.5f; + + // y: top --> down; x: left --> right + pts[0].x = box.x_ctr + sinTheta2 * box.h + cosTheta2 * box.w; + pts[0].y = box.y_ctr + cosTheta2 * box.h - sinTheta2 * box.w; + pts[1].x = box.x_ctr - sinTheta2 * box.h + cosTheta2 * box.w; + pts[1].y = box.y_ctr - cosTheta2 * box.h - sinTheta2 * box.w; + pts[2].x = 2 * box.x_ctr - pts[0].x; + pts[2].y = 2 * box.y_ctr - pts[0].y; + pts[3].x = 2 * box.x_ctr - pts[1].x; + pts[3].y = 2 * box.y_ctr - pts[1].y; +} + +template +inline int get_intersection_points( + const Point (&pts1)[4], + const Point (&pts2)[4], + Point (&intersections)[24]) { + // Line vector + // A line from p1 to p2 is: p1 + (p2-p1)*t, t=[0,1] + Point vec1[4], vec2[4]; + for (int i = 0; i < 4; i++) { + vec1[i] = pts1[(i + 1) % 4] - pts1[i]; + vec2[i] = pts2[(i + 1) % 4] - pts2[i]; + } + + // When computing the intersection area, it doesn't hurt if we have + // more (duplicated/approximate) intersections/vertices than needed, + // while it can cause drastic difference if we miss an intersection/vertex. + // Therefore, we add an epsilon to relax the comparisons between + // the float point numbers that decide the intersection points. + double EPS = 1e-5; + + // Line test - test all line combos for intersection + int num = 0; // number of intersections + for (int i = 0; i < 4; i++) { + for (int j = 0; j < 4; j++) { + // Solve for 2x2 Ax=b + T det = cross_2d(vec2[j], vec1[i]); + + // This takes care of parallel lines + if (fabs(det) <= 1e-14) { + continue; + } + + auto vec12 = pts2[j] - pts1[i]; + + T t1 = cross_2d(vec2[j], vec12) / det; + T t2 = cross_2d(vec1[i], vec12) / det; + + if (t1 > -EPS && t1 < 1.0f + EPS && t2 > -EPS && t2 < 1.0f + EPS) { + intersections[num++] = pts1[i] + vec1[i] * t1; + } + } + } + + // Check for vertices of rect1 inside rect2 + { + const auto& AB = vec2[0]; + const auto& DA = vec2[3]; + auto ABdotAB = dot_2d(AB, AB); + auto ADdotAD = dot_2d(DA, DA); + for (int i = 0; i < 4; i++) { + // assume ABCD is the rectangle, and P is the point to be judged + // P is inside ABCD iff. P's projection on AB lies within AB + // and P's projection on AD lies within AD + auto AP = pts1[i] - pts2[0]; + + auto APdotAB = dot_2d(AP, AB); + auto APdotAD = -dot_2d(AP, DA); + + if ((APdotAB > -EPS) && (APdotAD > -EPS) && (APdotAB < ABdotAB + EPS) && + (APdotAD < ADdotAD + EPS)) { + intersections[num++] = pts1[i]; + } + } + } + + // Reverse the check - check for vertices of rect2 inside rect1 + { + const auto& AB = vec1[0]; + const auto& DA = vec1[3]; + auto ABdotAB = dot_2d(AB, AB); + auto ADdotAD = dot_2d(DA, DA); + for (int i = 0; i < 4; i++) { + auto AP = pts2[i] - pts1[0]; + + auto APdotAB = dot_2d(AP, AB); + auto APdotAD = -dot_2d(AP, DA); + + if ((APdotAB > -EPS) && (APdotAD > -EPS) && (APdotAB < ABdotAB + EPS) && + (APdotAD < ADdotAD + EPS)) { + intersections[num++] = pts2[i]; + } + } + } + + return num; +} + +template +inline int convex_hull_graham( + const Point (&p)[24], + const int& num_in, + Point (&q)[24], + bool shift_to_zero = false) { + assert(num_in >= 2); + + // Step 1: + // Find point with minimum y + // if more than 1 points have the same minimum y, + // pick the one with the minimum x. + int t = 0; + for (int i = 1; i < num_in; i++) { + if (p[i].y < p[t].y || (p[i].y == p[t].y && p[i].x < p[t].x)) { + t = i; + } + } + auto& start = p[t]; // starting point + + // Step 2: + // Subtract starting point from every points (for sorting in the next step) + for (int i = 0; i < num_in; i++) { + q[i] = p[i] - start; + } + + // Swap the starting point to position 0 + auto tmp = q[0]; + q[0] = q[t]; + q[t] = tmp; + + // Step 3: + // Sort point 1 ~ num_in according to their relative cross-product values + // (essentially sorting according to angles) + // If the angles are the same, sort according to their distance to origin + T dist[24]; + for (int i = 0; i < num_in; i++) { + dist[i] = dot_2d(q[i], q[i]); + } + + for (int i = 1; i < num_in - 1; i++) { + for (int j = i + 1; j < num_in; j++) { + T crossProduct = cross_2d(q[i], q[j]); + if ((crossProduct < -1e-6) || + (fabs(crossProduct) < 1e-6 && dist[i] > dist[j])) { + auto q_tmp = q[i]; + q[i] = q[j]; + q[j] = q_tmp; + auto dist_tmp = dist[i]; + dist[i] = dist[j]; + dist[j] = dist_tmp; + } + } + } + + // compute distance to origin after sort, since the points are now different. + for (int i = 0; i < num_in; i++) { + dist[i] = dot_2d(q[i], q[i]); + } + + // Step 4: + // Make sure there are at least 2 points (that don't overlap with each other) + // in the stack + int k; // index of the non-overlapped second point + for (k = 1; k < num_in; k++) { + if (dist[k] > 1e-8) { + break; + } + } + if (k == num_in) { + // We reach the end, which means the convex hull is just one point + q[0] = p[t]; + return 1; + } + q[1] = q[k]; + int m = 2; // 2 points in the stack + + // Step 5: + // Finally we can start the scanning process. + // When a non-convex relationship between the 3 points is found + // (either concave shape or duplicated points), + // we pop the previous point from the stack + // until the 3-point relationship is convex again, or + // until the stack only contains two points + for (int i = k + 1; i < num_in; i++) { + while (m > 1) { + auto q1 = q[i] - q[m - 2], q2 = q[m - 1] - q[m - 2]; + // cross_2d() uses FMA and therefore computes round(round(q1.x*q2.y) - + // q2.x*q1.y) So it may not return 0 even when q1==q2. Therefore we + // compare round(q1.x*q2.y) and round(q2.x*q1.y) directly. (round means + // round to nearest floating point). + if (q1.x * q2.y >= q2.x * q1.y) { + m--; + } else { + break; + } + } + q[m++] = q[i]; + } + + // Step 6 (Optional): + // In general sense we need the original coordinates, so we + // need to shift the points back (reverting Step 2) + // But if we're only interested in getting the area/perimeter of the shape + // We can simply return. + if (!shift_to_zero) { + for (int i = 0; i < m; i++) { + q[i] += start; + } + } + + return m; +} + +template +inline T polygon_area(const Point (&q)[24], const int& m) { + if (m <= 2) { + return 0; + } + + T area = 0; + for (int i = 1; i < m - 1; i++) { + area += fabs(cross_2d(q[i] - q[0], q[i + 1] - q[0])); + } + + return area / 2.0; +} + +template +inline T rotated_boxes_intersection( + const RotatedBox& box1, + const RotatedBox& box2) { + // There are up to 4 x 4 + 4 + 4 = 24 intersections (including dups) returned + // from get_intersection_points + Point intersectPts[24], orderedPts[24]; + + Point pts1[4]; + Point pts2[4]; + get_rotated_vertices(box1, pts1); + get_rotated_vertices(box2, pts2); + + int num = get_intersection_points(pts1, pts2, intersectPts); + + if (num <= 2) { + return 0.0; + } + + // Convex Hull to order the intersection points in clockwise order and find + // the contour area. + int num_convex = convex_hull_graham(intersectPts, num, orderedPts, true); + return polygon_area(orderedPts, num_convex); +} + +template +inline T single_box_iou_rotated( + T const* const box1_raw, + T const* const box2_raw) { + // shift center to the middle point to achieve higher precision in result + RotatedBox box1, box2; + auto center_shift_x = (box1_raw[0] + box2_raw[0]) / 2.0; + auto center_shift_y = (box1_raw[1] + box2_raw[1]) / 2.0; + box1.x_ctr = box1_raw[0] - center_shift_x; + box1.y_ctr = box1_raw[1] - center_shift_y; + box1.w = box1_raw[2]; + box1.h = box1_raw[3]; + box1.a = box1_raw[4]; + box2.x_ctr = box2_raw[0] - center_shift_x; + box2.y_ctr = box2_raw[1] - center_shift_y; + box2.w = box2_raw[2]; + box2.h = box2_raw[3]; + box2.a = box2_raw[4]; + + T area1 = box1.w * box1.h; + T area2 = box2.w * box2.h; + if (area1 < 1e-14 || area2 < 1e-14) { + return 0.f; + } + + T intersection = rotated_boxes_intersection(box1, box2); + T iou = intersection / (area1 + area2 - intersection); + return iou; +} + +template +void box_iou_rotated_cpu_kernel( + const at::Tensor& boxes1, + const at::Tensor& boxes2, + at::Tensor& ious) { + auto num_boxes1 = boxes1.size(0); + auto num_boxes2 = boxes2.size(0); + + // Use accessors for efficient element access + auto boxes1_a = boxes1.accessor(); + auto boxes2_a = boxes2.accessor(); + auto ious_a = ious.accessor(); + + for (int64_t i = 0; i < num_boxes1; i++) { + for (int64_t j = 0; j < num_boxes2; j++) { + ious_a[i * num_boxes2 + j] = + single_box_iou_rotated(&boxes1_a[i][0], &boxes2_a[j][0]); + } + } +} + +at::Tensor box_iou_rotated_kernel( + const at::Tensor& boxes1, + const at::Tensor& boxes2) { + TORCH_CHECK(boxes1.is_cpu(), "boxes1 must be a CPU tensor"); + TORCH_CHECK(boxes2.is_cpu(), "boxes2 must be a CPU tensor"); + TORCH_CHECK( + boxes1.dim() == 2 && boxes1.size(1) == 5, + "boxes1 should have shape (N, 5), got ", + boxes1.sizes()); + TORCH_CHECK( + boxes2.dim() == 2 && boxes2.size(1) == 5, + "boxes2 should have shape (M, 5), got ", + boxes2.sizes()); + TORCH_CHECK( + boxes1.scalar_type() == boxes2.scalar_type(), + "boxes1 and boxes2 must have the same dtype"); + + auto num_boxes1 = boxes1.size(0); + auto num_boxes2 = boxes2.size(0); + + if (num_boxes1 == 0 || num_boxes2 == 0) { + return at::empty( + {num_boxes1, num_boxes2}, boxes1.options().dtype(at::kFloat)); + } + + auto boxes1_contiguous = boxes1.contiguous(); + auto boxes2_contiguous = boxes2.contiguous(); + + at::Tensor ious = + at::empty({num_boxes1 * num_boxes2}, boxes1.options().dtype(at::kFloat)); + + AT_DISPATCH_FLOATING_TYPES(boxes1.scalar_type(), "box_iou_rotated_cpu", [&] { + box_iou_rotated_cpu_kernel( + boxes1_contiguous, boxes2_contiguous, ious); + }); + + return ious.reshape({num_boxes1, num_boxes2}); +} + +} // namespace + +TORCH_LIBRARY_IMPL(torchvision, CPU, m) { + m.impl( + TORCH_SELECTIVE_NAME("torchvision::box_iou_rotated"), + TORCH_FN(box_iou_rotated_kernel)); +} + +} // namespace ops +} // namespace vision diff --git a/csrc/src/exports.cpp b/csrc/src/exports.cpp index 4c1dd05..6bfd586 100644 --- a/csrc/src/exports.cpp +++ b/csrc/src/exports.cpp @@ -21,6 +21,13 @@ TORCHVISIONLIB_API void* _vision_ops_nms (void* dets, void* scores, double iou_t } TORCHVISIONLIB_HANDLE_EXCEPTION return (void*) NULL; } +torch::Tensor vision_ops_box_iou_rotated (torch::Tensor boxes1, torch::Tensor boxes2); +TORCHVISIONLIB_API void* _vision_ops_box_iou_rotated (void* boxes1, void* boxes2) { + try { + return make_raw::Tensor(vision_ops_box_iou_rotated(from_raw::Tensor(boxes1), from_raw::Tensor(boxes2))); + } 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 46439fa..c789d00 100644 --- a/csrc/src/ops.cpp +++ b/csrc/src/ops.cpp @@ -7,6 +7,7 @@ #include #include #include +#include "ops/box_iou_rotated/box_iou_rotated.h" // [[torch::export]] @@ -14,6 +15,11 @@ torch::Tensor vision_ops_nms(torch::Tensor dets, torch::Tensor scores, double io return vision::ops::nms(dets, scores, iou_threshold); } +// [[torch::export]] +torch::Tensor vision_ops_box_iou_rotated(torch::Tensor boxes1, torch::Tensor boxes2) { + return vision::ops::box_iou_rotated(boxes1, boxes2); +} + // [[torch::export]] torch::Tensor vision_ops_deform_conv2d( torch::Tensor input, diff --git a/csrc/src/torchvisionlib.def b/csrc/src/torchvisionlib.def index df913d4..26a0602 100644 --- a/csrc/src/torchvisionlib.def +++ b/csrc/src/torchvisionlib.def @@ -4,6 +4,7 @@ EXPORTS ;------ autogenerated ------------------------- ; don't modify between the autogenerated lines _vision_ops_nms + _vision_ops_box_iou_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 df913d4..26a0602 100644 --- a/inst/def/torchvisionlib.def +++ b/inst/def/torchvisionlib.def @@ -4,6 +4,7 @@ EXPORTS ;------ autogenerated ------------------------- ; don't modify between the autogenerated lines _vision_ops_nms + _vision_ops_box_iou_rotated _vision_ops_deform_conv2d _vision_ops_ps_roi_align _vision_ops_ps_roi_pool diff --git a/inst/include/torchvisionlib/exports.h b/inst/include/torchvisionlib/exports.h index a8de27c..47f42da 100644 --- a/inst/include/torchvisionlib/exports.h +++ b/inst/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(); diff --git a/man/ops_box_iou_rotated.Rd b/man/ops_box_iou_rotated.Rd new file mode 100644 index 0000000..11d3b18 --- /dev/null +++ b/man/ops_box_iou_rotated.Rd @@ -0,0 +1,40 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/ops.R +\name{ops_box_iou_rotated} +\alias{ops_box_iou_rotated} +\title{Intersection-over-union between rotated boxes} +\usage{ +ops_box_iou_rotated(boxes1, boxes2, fmt = "cxcywhr") +} +\arguments{ +\item{boxes1}{\code{Tensor[N, K]} first set of rotated boxes.} + +\item{boxes2}{\code{Tensor[M, K]} second set of rotated boxes.} + +\item{fmt}{format of the input boxes. One of: +\itemize{ +\item \code{"cxcywhr"} (\code{K = 5}): center \verb{(cx, cy)}, width, height and rotation +angle \code{r} in degrees (counter-clockwise positive). +\item \code{"xywhr"} (\code{K = 5}): top-left corner \verb{(x1, y1)}, width, height and angle. +\item \code{"xyxyxyxy"} (\code{K = 8}): the four corners +\verb{(x1, y1, x2, y2, x3, y3, x4, y4)}. +}} +} +\value{ +\code{Tensor[N, M]} float32 matrix of pairwise IoU values. +} +\description{ +Computes the pairwise intersection-over-union (IoU) between two sets of +rotated bounding boxes. +} +\examples{ +if (torchvisionlib_is_installed()) { + boxes <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 45), nrow = 1)) + ops_box_iou_rotated(boxes, boxes) +} +} +\seealso{ +Other ops: +\code{\link[=ops_nms]{ops_nms()}} +} +\concept{ops} diff --git a/man/ops_nms.Rd b/man/ops_nms.Rd index 3110dd9..f612726 100644 --- a/man/ops_nms.Rd +++ b/man/ops_nms.Rd @@ -36,4 +36,8 @@ if (torchvisionlib_is_installed()) { ops_nms(torch::torch_rand(3, 4), torch::torch_rand(3), 0.5) } } +\seealso{ +Other ops: +\code{\link[=ops_box_iou_rotated]{ops_box_iou_rotated()}} +} \concept{ops} diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index 4ee93d8..3ccaa44 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -24,6 +24,18 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// rcpp_vision_ops_box_iou_rotated +torch::Tensor rcpp_vision_ops_box_iou_rotated(torch::Tensor boxes1, torch::Tensor boxes2); +RcppExport SEXP _torchvisionlib_rcpp_vision_ops_box_iou_rotated(SEXP boxes1SEXP, SEXP boxes2SEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< torch::Tensor >::type boxes1(boxes1SEXP); + Rcpp::traits::input_parameter< torch::Tensor >::type boxes2(boxes2SEXP); + rcpp_result_gen = Rcpp::wrap(rcpp_vision_ops_box_iou_rotated(boxes1, boxes2)); + 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) { @@ -168,6 +180,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_box_iou_rotated", (DL_FUNC) &_torchvisionlib_rcpp_vision_ops_box_iou_rotated, 2}, {"_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 142362a..de94f47 100644 --- a/src/exports.cpp +++ b/src/exports.cpp @@ -8,6 +8,10 @@ torch::Tensor rcpp_vision_ops_nms (torch::Tensor dets, torch::Tensor scores, dou return vision_ops_nms(dets.get(), scores.get(), iou_threshold); } // [[Rcpp::export]] +torch::Tensor rcpp_vision_ops_box_iou_rotated (torch::Tensor boxes1, torch::Tensor boxes2) { + return vision_ops_box_iou_rotated(boxes1.get(), boxes2.get()); +} +// [[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 534858c..a259e9d 100644 --- a/src/exports.h +++ b/src/exports.h @@ -4,6 +4,7 @@ #include "torchvisionlib_types.h" torch::Tensor rcpp_vision_ops_nms (torch::Tensor dets, torch::Tensor scores, double iou_threshold); +torch::Tensor rcpp_vision_ops_box_iou_rotated (torch::Tensor boxes1, torch::Tensor boxes2); 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/testthat/test-ops-box-iou-rotated.R b/tests/testthat/test-ops-box-iou-rotated.R new file mode 100644 index 0000000..97df2bf --- /dev/null +++ b/tests/testthat/test-ops-box-iou-rotated.R @@ -0,0 +1,167 @@ +dtypes <- list(torch::torch_float32(), torch::torch_float64()) + +# cxcywhr boxes used by the core test, together with their expected pairwise +# IoU. Squares 1-4 are the same box rotated in 90-degree steps (IoU 1); box 5 +# is far away (IoU 0); the 45/135-degree rectangles 6-8 overlap in a cross +# (IoU 1/3) or coincide (IoU 1). Ported from torchvision's TestRotatedBoxIou. +rotated_boxes <- torch::torch_tensor(matrix(c( + 0, 0, 10, 10, 45, + 0, 0, 10, 10, 135, + 0, 0, 10, 10, -45, + 0, 0, 10, 10, -135, + 100, 100, 10, 10, 30, + 50, 50, 20, 10, 45, + 50, 50, 20, 10, 135, + 50, 50, 20, 10, -135 +), ncol = 5, byrow = TRUE)) + +rotated_boxes_expected <- torch::torch_tensor(matrix(c( + 1, 1, 1, 1, 0, 0, 0, 0, + 1, 1, 1, 1, 0, 0, 0, 0, + 1, 1, 1, 1, 0, 0, 0, 0, + 1, 1, 1, 1, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 1/3, 1, + 0, 0, 0, 0, 0, 1/3, 1, 1/3, + 0, 0, 0, 0, 0, 1, 1/3, 1 +), ncol = 8, byrow = TRUE), dtype = torch::torch_float32()) + +# cxcywhr -> xywhr / xyxyxyxy, used to check the R-side format conversions +# against the cxcywhr path. Formulas mirror torchvision's box_convert. +cxcywhr_to_xywhr <- function(boxes) { + b <- torch::torch_unbind(boxes, dim = -1) + cx <- b[[1]] + cy <- b[[2]] + w <- b[[3]] + h <- b[[4]] + r <- b[[5]] + rad <- r * pi / 180 + cos <- torch::torch_cos(rad) + sin <- torch::torch_sin(rad) + x1 <- cx - w / 2 * cos - h / 2 * sin + y1 <- cy - h / 2 * cos + w / 2 * sin + torch::torch_stack(list(x1, y1, w, h, r), dim = -1) +} + +xywhr_to_xyxyxyxy <- function(boxes) { + b <- torch::torch_unbind(boxes, dim = -1) + x1 <- b[[1]] + y1 <- b[[2]] + w <- b[[3]] + h <- b[[4]] + r <- b[[5]] + rad <- r * pi / 180 + cos <- torch::torch_cos(rad) + sin <- torch::torch_sin(rad) + x2 <- x1 + w * cos + y2 <- y1 - w * sin + x3 <- x2 + h * sin + y3 <- y2 + h * cos + x4 <- x1 + h * sin + y4 <- y1 + h * cos + torch::torch_stack(list(x1, y1, x2, y2, x3, y3, x4, y4), dim = -1) +} + +test_that("box_iou_rotated matches expected IoU for cxcywhr boxes", { + for (dtype in dtypes) { + boxes <- rotated_boxes$to(dtype = dtype) + out <- ops_box_iou_rotated(boxes, boxes) + # The op returns float32 regardless of the input dtype. + expect_true(out$dtype == torch::torch_float32()) + expect_equal_to_tensor(out, rotated_boxes_expected, atol = 1e-4, rtol = 1e-4) + } +}) + +test_that("box_iou_rotated is invariant to angle sign and 180-degree rotation", { + for (dtype in dtypes) { + base <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 0), nrow = 1), dtype = dtype) + for (angle in c(45, 53, 195)) { + pos <- torch::torch_tensor(matrix(c(0, 0, 10, 10, angle), nrow = 1), dtype = dtype) + neg <- torch::torch_tensor(matrix(c(0, 0, 10, 10, -angle), nrow = 1), dtype = dtype) + expect_equal_to_tensor( + ops_box_iou_rotated(base, pos), ops_box_iou_rotated(base, neg) + ) + } + + rect <- torch::torch_tensor(matrix(c(50, 50, 20, 10, 30), nrow = 1), dtype = dtype) + rect180 <- torch::torch_tensor(matrix(c(50, 50, 20, 10, 30 - 180), nrow = 1), dtype = dtype) + expect_equal(as.numeric(ops_box_iou_rotated(rect, rect180)), 1, tolerance = 1e-5) + } +}) + +# Deliberately float64: recovering the angle from corners via atan2 loses enough +# float32 precision that upstream needs atol = 0.5 for xyxyxyxy on macOS. Using +# float64 keeps this a tight check of the conversions rather than of float noise. +test_that("box_iou_rotated gives the same result across formats", { + boxes <- rotated_boxes$to(dtype = torch::torch_float64()) + ref <- ops_box_iou_rotated(boxes, boxes, fmt = "cxcywhr") + + xywhr <- cxcywhr_to_xywhr(boxes) + expect_equal_to_tensor( + ops_box_iou_rotated(xywhr, xywhr, fmt = "xywhr"), ref, atol = 1e-4 + ) + + xyxyxyxy <- xywhr_to_xyxyxyxy(xywhr) + expect_equal_to_tensor( + ops_box_iou_rotated(xyxyxyxy, xyxyxyxy, fmt = "xyxyxyxy"), ref, atol = 1e-4 + ) +}) + +test_that("box_iou_rotated handles containment and zero-area boxes", { + for (dtype in dtypes) { + outer <- torch::torch_tensor(matrix(c(0, 0, 20, 20, 45), nrow = 1), dtype = dtype) + inner <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 45), nrow = 1), dtype = dtype) + expect_equal(as.numeric(ops_box_iou_rotated(outer, inner)), 0.25, tolerance = 1e-5) + + degenerate <- torch::torch_tensor(matrix(c(0, 0, 0, 10, 45), nrow = 1), dtype = dtype) + other <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 30), nrow = 1), dtype = dtype) + expect_equal(as.numeric(ops_box_iou_rotated(degenerate, other)), 0) + } +}) + +test_that("box_iou_rotated handles empty inputs and output shape", { + for (dtype in dtypes) { + boxes <- torch::torch_zeros(c(10, 5), dtype = dtype) + empty <- torch::torch_zeros(c(0, 5), dtype = dtype) + expect_equal(ops_box_iou_rotated(empty, boxes)$shape, c(0, 10)) + expect_equal(ops_box_iou_rotated(boxes, empty)$shape, c(10, 0)) + + b1 <- torch::torch_rand(5, 5)$to(dtype = dtype) + b2 <- torch::torch_rand(7, 5)$to(dtype = dtype) + expect_equal(ops_box_iou_rotated(b1, b2)$shape, c(5, 7)) + } +}) + +test_that("box_iou_rotated errors on an unknown format", { + boxes <- torch::torch_tensor(matrix(c(0, 0, 10, 10, 0), nrow = 1)) + expect_error( + ops_box_iou_rotated(boxes, boxes, fmt = "nope"), + regexp = "Unsupported format" + ) +}) + +test_that("box_iou_rotated is numerically stable (Detectron2 regressions)", { + for (dtype in dtypes) { + # Precision at large coordinates: IoU is the height ratio. + b1 <- torch::torch_tensor(matrix(c(565, 565, 10, 10, 0), nrow = 1), dtype = dtype) + b2 <- torch::torch_tensor(matrix(c(565, 565, 10, 8.3, 0), nrow = 1), dtype = dtype) + expect_equal(as.numeric(ops_box_iou_rotated(b1, b2)), 8.3 / 10, tolerance = 1e-5) + + # Nearly identical large boxes should have IoU close to 1. + b3 <- torch::torch_tensor(matrix(c(2563.74462890625, 1436.7901611328125, + 2174.703369140625, 214.095001220703125, + 115.11834716796875), nrow = 1), dtype = dtype) + b4 <- torch::torch_tensor(matrix(c(2563.74462890625, 1436.790283203125, + 2174.702880859375, 214.0949554443359375, + 115.11835479736328125), nrow = 1), dtype = dtype) + expect_equal(as.numeric(ops_box_iou_rotated(b3, b4)), 1, tolerance = 1e-5) + + # Extreme coordinates must not push the IoU outside [0, 1]. + b5 <- torch::torch_tensor(matrix(c(160, 153, 230, 23, -37), nrow = 1), dtype = dtype) + b6 <- torch::torch_tensor(matrix(c(-1.117407639806935e17, 1.3858420478349148e18, + 1000, 1000, 1612), nrow = 1), dtype = dtype) + iou <- as.numeric(ops_box_iou_rotated(b5, b6)) + expect_gte(iou, 0) + expect_lte(iou, 1) + } +})