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