From bb4c65d8ffcfb64fa75cfc8ad3c1b66f5040a6b4 Mon Sep 17 00:00:00 2001
From: Range King <RangeKingHZ@gmail.com>
Date: Fri, 27 May 2022 10:21:48 +0800
Subject: [PATCH] Add type hints in bbox.py and box_iou_rotated.py (#1981)

* Add type hints in bbox.py

* Add type hints in box_iou_rotated.py

* Update ball_query.py

* Revert "Add type hints in bbox.py"

This reverts commit 102566488e3aa9328b1c6e1159e2699d40f9f9f5.

* Update bbox.py

* Revert "Update bbox.py"

This reverts commit d34b2c6ed54ba024b9674c948493001a8b4ff864.

* Update bbox.py
---
 mmcv/ops/ball_query.py      |  2 +-
 mmcv/ops/bbox.py            | 12 ++++++++++--
 mmcv/ops/box_iou_rotated.py | 12 +++++++-----
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/mmcv/ops/ball_query.py b/mmcv/ops/ball_query.py
index f296b5463..61a5a5c90 100644
--- a/mmcv/ops/ball_query.py
+++ b/mmcv/ops/ball_query.py
@@ -18,7 +18,7 @@ class BallQuery(Function):
             min_radius (float): minimum radius of the balls.
             max_radius (float): maximum radius of the balls.
             sample_num (int): maximum number of features in the balls.
-            xyz (Tensor): (B, N, 3) xyz coordinates of the features.
+            xyz (torch.Tensor): (B, N, 3) xyz coordinates of the features.
             center_xyz (torch.Tensor): (B, npoint, 3) centers of the ball
                 query.
 
diff --git a/mmcv/ops/bbox.py b/mmcv/ops/bbox.py
index e28cd0a91..bf6bd43bb 100644
--- a/mmcv/ops/bbox.py
+++ b/mmcv/ops/bbox.py
@@ -6,7 +6,11 @@ from ..utils import ext_loader
 ext_module = ext_loader.load_ext('_ext', ['bbox_overlaps'])
 
 
-def _bbox_overlaps_cpu(bboxes1, bboxes2, mode='iou', aligned=False, offset=0):
+def _bbox_overlaps_cpu(bboxes1: torch.Tensor,
+                       bboxes2: torch.Tensor,
+                       mode: str = 'iou',
+                       aligned: bool = False,
+                       offset: int = 0) -> torch.Tensor:
     assert mode in ['iou', 'iof']
 
     if aligned:
@@ -43,7 +47,11 @@ def _bbox_overlaps_cpu(bboxes1, bboxes2, mode='iou', aligned=False, offset=0):
     return ious
 
 
-def bbox_overlaps(bboxes1, bboxes2, mode='iou', aligned=False, offset=0):
+def bbox_overlaps(bboxes1: torch.Tensor,
+                  bboxes2: torch.Tensor,
+                  mode: str = 'iou',
+                  aligned: bool = False,
+                  offset: int = 0) -> torch.Tensor:
     """Calculate overlap between two set of bboxes.
 
     If ``aligned`` is ``False``, then calculate the ious between each bbox
diff --git a/mmcv/ops/box_iou_rotated.py b/mmcv/ops/box_iou_rotated.py
index 20f663959..2443af27c 100644
--- a/mmcv/ops/box_iou_rotated.py
+++ b/mmcv/ops/box_iou_rotated.py
@@ -1,14 +1,16 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import torch
+
 from ..utils import ext_loader
 
 ext_module = ext_loader.load_ext('_ext', ['box_iou_rotated'])
 
 
-def box_iou_rotated(bboxes1,
-                    bboxes2,
-                    mode='iou',
-                    aligned=False,
-                    clockwise=True):
+def box_iou_rotated(bboxes1: torch.Tensor,
+                    bboxes2: torch.Tensor,
+                    mode: str = 'iou',
+                    aligned: bool = False,
+                    clockwise: bool = True) -> torch.Tensor:
     """Return intersection-over-union (Jaccard index) of boxes.
 
     Both sets of boxes are expected to be in