From 2045abd954c61c8089f1b7e42490561059ebbfe6 Mon Sep 17 00:00:00 2001
From: hust17yixuan <303660421@qq.com>
Date: Fri, 24 Jan 2025 09:45:53 +0800
Subject: [PATCH] Revert "adapt boxes_overlap_bev & box_iou_rotated (#3162)"

This reverts commit 44d871e7f36bb00afc3946ea36cc6dcf03e07e44.
---
 mmcv/ops/box_iou_rotated.py                   |  5 +++
 .../csrc/pytorch/npu/box_iou_rotated_npu.cpp  | 37 +++++++++++++++++--
 tests/test_ops/test_iou3d.py                  |  6 +--
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/mmcv/ops/box_iou_rotated.py b/mmcv/ops/box_iou_rotated.py
index 8e199d9ac8..a811531d42 100644
--- a/mmcv/ops/box_iou_rotated.py
+++ b/mmcv/ops/box_iou_rotated.py
@@ -142,6 +142,11 @@ def box_iou_rotated(bboxes1: torch.Tensor,
         flip_mat[-1] = -1
         bboxes1 = bboxes1 * flip_mat
         bboxes2 = bboxes2 * flip_mat
+    if bboxes1.device.type == 'npu':
+        scale_mat = bboxes1.new_ones(bboxes1.shape[-1])
+        scale_mat[-1] = 1.0 / 0.01745329252
+        bboxes1 = bboxes1 * scale_mat
+        bboxes2 = bboxes2 * scale_mat
     bboxes1 = bboxes1.contiguous()
     bboxes2 = bboxes2.contiguous()
     ext_module.box_iou_rotated(
diff --git a/mmcv/ops/csrc/pytorch/npu/box_iou_rotated_npu.cpp b/mmcv/ops/csrc/pytorch/npu/box_iou_rotated_npu.cpp
index fd68af68d9..c6e6b66478 100644
--- a/mmcv/ops/csrc/pytorch/npu/box_iou_rotated_npu.cpp
+++ b/mmcv/ops/csrc/pytorch/npu/box_iou_rotated_npu.cpp
@@ -8,11 +8,40 @@ void box_iou_rotated_impl(const Tensor boxes1, const Tensor boxes2, Tensor ious,
 
 void box_iou_rotated_npu(const Tensor boxes1, const Tensor boxes2, Tensor ious,
                          const int mode_flag, const bool aligned) {
-  TORCH_CHECK(boxes1.size(1) == 5, "boxes1 must be 2D tensor (N, 5)");
-  TORCH_CHECK(boxes1.size(1) == 5, "boxes1 must be 2D tensor (N, 5)");
+  at::Tensor boxes = at::ones_like(boxes1);
+  at::Tensor query_boxes = at::ones_like(boxes2);
+  boxes = boxes1.transpose(0, 1).unsqueeze(0);
+  query_boxes = boxes2.transpose(0, 1).unsqueeze(0);
 
-  EXEC_NPU_CMD(aclnnBoxIou, boxes1, boxes2, mode_flag, aligned, ious);
-  return;
+  bool is_trans = false;
+  string modeStr = "iou";
+  if (mode_flag == 1) {
+    modeStr = "iof";
+  }
+  bool is_cross = true;
+  if (aligned) {
+    is_cross = false;
+  }
+  float v_threshold = 0;
+  float e_threshold = 0;
+
+  OpCommand cmd;
+  cmd.Name("RotatedIou")
+      .Input(boxes)
+      .Input(query_boxes)
+      .Output(ious)
+      .Attr("trans", is_trans)
+      .Attr("mode", modeStr)
+      .Attr("is_cross", is_cross)
+      .Attr("v_threshold", v_threshold)
+      .Attr("e_threshold", e_threshold)
+      .Run();
+
+  if (is_cross) {
+    ious = ious.view({boxes1.size(0), boxes2.size(0)});
+  } else {
+    ious = ious.view({boxes1.size(0), 1});
+  }
 }
 
 REGISTER_NPU_IMPL(box_iou_rotated_impl, box_iou_rotated_npu);
diff --git a/tests/test_ops/test_iou3d.py b/tests/test_ops/test_iou3d.py
index 6b2456e8b9..27a09eb361 100644
--- a/tests/test_ops/test_iou3d.py
+++ b/tests/test_ops/test_iou3d.py
@@ -11,11 +11,7 @@
     pytest.param(
         'cuda',
         marks=pytest.mark.skipif(
-            not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
-    pytest.param(
-        'npu',
-        marks=pytest.mark.skipif(
-            not IS_NPU_AVAILABLE, reason='requires NPU support'))
+            not IS_CUDA_AVAILABLE, reason='requires CUDA support'))
 ])
 def test_boxes_overlap_bev(device):
     np_boxes1 = np.asarray([[1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 0.0],