mirrors
/
mmcv
mirror of https://github.com/open-mmlab/mmcv.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Add type hints in bbox.py

pull/1981/head
RangeKing 2022-05-25 18:55:19 +08:00
parent c90f2be0be
commit 102566488e
1 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mmcv/ops/bbox.py
View File

@ -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