[Feature] Add task_modules type hint

mmdet3d/models/task_modules/anchor/anchor_3d_generator.py

@@ -1,6 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Tuple, Union
+
 import mmengine
+import numpy as np
 import torch
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 
@@ -37,13 +41,13 @@ class Anchor3DRangeGenerator(object):
  """
 
  def __init__(self,
- ranges,
- sizes=[[3.9, 1.6, 1.56]],
- scales=[1],
- rotations=[0, 1.5707963],
- custom_values=(),
- reshape_out=True,
- size_per_range=True):
+ ranges: List[List[float]],
+ sizes: List[List[float]] = [[3.9, 1.6, 1.56]],
+ scales: List[int] = [1],
+ rotations: List[float] = [0, 1.5707963],
+ custom_values: Tuple[float] = (),
+ reshape_out: bool = True,
+ size_per_range: bool = True) -> None:
  assert mmengine.is_list_of(ranges, list)
  if size_per_range:
  if len(sizes) != len(ranges):
@@ -64,7 +68,7 @@ def __init__(self,
  self.reshape_out = reshape_out
  self.size_per_range = size_per_range
 
- def __repr__(self):
+ def __repr__(self) -> str:
  s = self.__class__.__name__ + '('
  s += f'anchor_range={self.ranges},\n'
  s += f'scales={self.scales},\n'
@@ -75,18 +79,20 @@ def __repr__(self):
  return s
 
  @property
- def num_base_anchors(self):
- """list[int]: Total number of base anchors in a feature grid."""
+ def num_base_anchors(self) -> int:
+ """int: Total number of base anchors in a feature grid."""
  num_rot = len(self.rotations)
  num_size = torch.tensor(self.sizes).reshape(-1, 3).size(0)
  return num_rot * num_size
 
  @property
- def num_levels(self):
+ def num_levels(self) -> int:
  """int: Number of feature levels that the generator is applied to."""
  return len(self.scales)
 
- def grid_anchors(self, featmap_sizes, device='cuda'):
+ def grid_anchors(self,
+ featmap_sizes: List[tuple],
+ device: str = 'cuda') -> List[Tensor]:
  """Generate grid anchors in multiple feature levels.
 
  Args:
@@ -112,7 +118,10 @@ def grid_anchors(self, featmap_sizes, device='cuda'):
  multi_level_anchors.append(anchors)
  return multi_level_anchors
 
- def single_level_grid_anchors(self, featmap_size, scale, device='cuda'):
+ def single_level_grid_anchors(self,
+ featmap_size: Tuple[int],
+ scale: float,
+ device: str = 'cuda') -> Tensor:
  """Generate grid anchors of a single level feature map.
 
  This function is usually called by method ``self.grid_anchors``.
@@ -153,12 +162,12 @@ def single_level_grid_anchors(self, featmap_size, scale, device='cuda'):
  return mr_anchors
 
  def anchors_single_range(self,
- feature_size,
- anchor_range,
- scale=1,
- sizes=[[3.9, 1.6, 1.56]],
+ feature_size: Union[List[float], Tuple[float]],
+ anchor_range: Union[Tensor, List[float]],
+ scale: Union[float, int] = 1,
+ sizes: Union[Tensor, int] = [[3.9, 1.6, 1.56]],
  rotations=[0, 1.5707963],
- device='cuda'):
+ device='cuda') -> Tensor:
  """Generate anchors in a single range.
 
  Args:
@@ -248,17 +257,19 @@ class AlignedAnchor3DRangeGenerator(Anchor3DRangeGenerator):
  center of the corresponding greature grid. Defaults to False.
  """
 
- def __init__(self, align_corner=False, **kwargs):
+ def __init__(self, align_corner: bool = False, **kwargs) -> None:
  super(AlignedAnchor3DRangeGenerator, self).__init__(**kwargs)
  self.align_corner = align_corner
 
  def anchors_single_range(self,
- feature_size,
- anchor_range,
- scale,
- sizes=[[3.9, 1.6, 1.56]],
- rotations=[0, 1.5707963],
- device='cuda'):
+ feature_size: Union[List[float], Tuple[float]],
+ anchor_range: Union[Tensor, List[float]],
+ scale: Union[float, int],
+ sizes: Union[List[list], np.ndarray,
+ Tensor] = [[3.9, 1.6, 1.56]],
+ rotations: Union[List[float], np.ndarray,
+ Tensor] = [0, 1.5707963],
+ device: str = 'cuda') -> Tensor:
  """Generate anchors in a single range.
 
  Args:
@@ -352,12 +363,14 @@ class AlignedAnchor3DRangeGeneratorPerCls(AlignedAnchor3DRangeGenerator):
  :class:`AlignedAnchor3DRangeGenerator`.
  """
 
- def __init__(self, **kwargs):
+ def __init__(self, **kwargs) -> None:
  super(AlignedAnchor3DRangeGeneratorPerCls, self).__init__(**kwargs)
  assert len(self.scales) == 1, 'Multi-scale feature map levels are' + \
  ' not supported currently in this kind of anchor generator.'
 
- def grid_anchors(self, featmap_sizes, device='cuda'):
+ def grid_anchors(self,
+ featmap_sizes: List[tuple],
+ device: str = 'cuda') -> List[List[Tensor]]:
  """Generate grid anchors in multiple feature levels.
 
  Args:
@@ -379,7 +392,10 @@ def grid_anchors(self, featmap_sizes, device='cuda'):
  multi_level_anchors.append(anchors)
  return multi_level_anchors
 
- def multi_cls_grid_anchors(self, featmap_sizes, scale, device='cuda'):
+ def multi_cls_grid_anchors(self,
+ featmap_sizes: List[tuple],
+ scale: float,
+ device: str = 'cuda') -> Tensor:
  """Generate grid anchors of a single level feature map for multi-class
  with different feature map sizes.
 

mmdet3d/models/task_modules/anchor/builder.py

@@ -1,21 +1,23 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
+from typing import Any
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.utils import ConfigType
 
 PRIOR_GENERATORS = TASK_UTILS
 
 ANCHOR_GENERATORS = TASK_UTILS
 
 
-def build_prior_generator(cfg, default_args=None):
+def build_prior_generator(cfg: ConfigType, default_args=None) -> Any:
  warnings.warn(
  '``build_prior_generator`` would be deprecated soon, please use '
  '``mmdet3d.registry.TASK_UTILS.build()`` ')
  return TASK_UTILS.build(cfg, default_args=default_args)
 
 
-def build_anchor_generator(cfg, default_args=None):
+def build_anchor_generator(cfg: ConfigType, default_args=None) -> Any:
  warnings.warn(
  '``build_anchor_generator`` would be deprecated soon, please use '
  '``mmdet3d.registry.TASK_UTILS.build()`` ')

mmdet3d/models/task_modules/assigners/max_3d_iou_assigner.py

@@ -42,16 +42,18 @@ class Max3DIoUAssigner(MaxIoUAssigner):
  iou_calculator (dict): Config of overlaps Calculator.
  """
 
- def __init__(self,
- pos_iou_thr: float,
- neg_iou_thr: Union[float, tuple],
- min_pos_iou: float = .0,
- gt_max_assign_all: bool = True,
- ignore_iof_thr: float = -1,
- ignore_wrt_candidates: bool = True,
- match_low_quality: bool = True,
- gpu_assign_thr: float = -1,
- iou_calculator: dict = dict(type='BboxOverlaps2D')):
+ def __init__(
+ self,
+ pos_iou_thr: float,
+ neg_iou_thr: Union[float, tuple],
+ min_pos_iou: float = .0,
+ gt_max_assign_all: bool = True,
+ ignore_iof_thr: float = -1,
+ ignore_wrt_candidates: bool = True,
+ match_low_quality: bool = True,
+ gpu_assign_thr: float = -1,
+ iou_calculator: dict = dict(type='BboxOverlaps2D')
+ ) -> None:
  self.pos_iou_thr = pos_iou_thr
  self.neg_iou_thr = neg_iou_thr
  self.min_pos_iou = min_pos_iou

mmdet3d/models/task_modules/builder.py

@@ -1,28 +1,30 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import warnings
+from typing import Any
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.utils.typing_utils import ConfigType
 
 BBOX_ASSIGNERS = TASK_UTILS
 BBOX_SAMPLERS = TASK_UTILS
 BBOX_CODERS = TASK_UTILS
 
 
-def build_assigner(cfg, **default_args):
+def build_assigner(cfg: ConfigType, **default_args) -> Any:
  """Builder of box assigner."""
  warnings.warn('``build_assigner`` would be deprecated soon, please use '
  '``mmdet3d.registry.TASK_UTILS.build()`` ')
  return TASK_UTILS.build(cfg, default_args=default_args)
 
 
-def build_sampler(cfg, **default_args):
+def build_sampler(cfg: ConfigType, **default_args) -> Any:
  """Builder of box sampler."""
  warnings.warn('``build_sampler`` would be deprecated soon, please use '
  '``mmdet3d.registry.TASK_UTILS.build()`` ')
  return TASK_UTILS.build(cfg, default_args=default_args)
 
 
-def build_bbox_coder(cfg, **default_args):
+def build_bbox_coder(cfg: ConfigType, **default_args) -> Any:
  """Builder of box coder."""
  warnings.warn('``build_bbox_coder`` would be deprecated soon, please use '
  '``mmdet3d.registry.TASK_UTILS.build()`` ')

mmdet3d/models/task_modules/coders/anchor_free_bbox_coder.py

@@ -1,8 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict
+
 import numpy as np
 import torch
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures import BaseInstance3DBoxes
 from .partial_bin_based_bbox_coder import PartialBinBasedBBoxCoder
 
 
@@ -15,13 +19,14 @@ class AnchorFreeBBoxCoder(PartialBinBasedBBoxCoder):
  with_rot (bool): Whether the bbox is with rotation.
  """
 
- def __init__(self, num_dir_bins, with_rot=True):
+ def __init__(self, num_dir_bins: int, with_rot: bool = True) -> None:
  super(AnchorFreeBBoxCoder, self).__init__(
  num_dir_bins, 0, [], with_rot=with_rot)
  self.num_dir_bins = num_dir_bins
  self.with_rot = with_rot
 
- def encode(self, gt_bboxes_3d, gt_labels_3d):
+ def encode(self, gt_bboxes_3d: BaseInstance3DBoxes,
+ gt_labels_3d: Tensor) -> tuple:
  """Encode ground truth to prediction targets.
 
  Args:
@@ -51,7 +56,7 @@ def encode(self, gt_bboxes_3d, gt_labels_3d):
  return (center_target, size_res_target, dir_class_target,
  dir_res_target)
 
- def decode(self, bbox_out):
+ def decode(self, bbox_out: dict) -> Tensor:
  """Decode predicted parts to bbox3d.
 
  Args:
@@ -85,7 +90,8 @@ def decode(self, bbox_out):
  bbox3d = torch.cat([center, bbox_size, dir_angle], dim=-1)
  return bbox3d
 
- def split_pred(self, cls_preds, reg_preds, base_xyz):
+ def split_pred(self, cls_preds: Tensor, reg_preds: Tensor,
+ base_xyz: Tensor) -> Dict[str, Tensor]:
  """Split predicted features to specific parts.
 
  Args:

mmdet3d/models/task_modules/coders/centerpoint_bbox_coders.py

@@ -1,6 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Tuple
+
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 
@@ -22,13 +25,13 @@ class CenterPointBBoxCoder(BaseBBoxCoder):
  """
 
  def __init__(self,
- pc_range,
- out_size_factor,
- voxel_size,
- post_center_range=None,
- max_num=100,
- score_threshold=None,
- code_size=9):
+ pc_range: List[float],
+ out_size_factor: int,
+ voxel_size: List[float],
+ post_center_range: List[float] = None,
+ max_num: int = 100,
+ score_threshold: float = None,
+ code_size: int = 9) -> None:
 
  self.pc_range = pc_range
  self.out_size_factor = out_size_factor
@@ -38,7 +41,10 @@ def __init__(self,
  self.score_threshold = score_threshold
  self.code_size = code_size
 
- def _gather_feat(self, feats, inds, feat_masks=None):
+ def _gather_feat(self,
+ feats: Tensor,
+ inds: Tensor,
+ feat_masks: Tensor = None) -> Tensor:
  """Given feats and indexes, returns the gathered feats.
 
  Args:
@@ -60,7 +66,7 @@ def _gather_feat(self, feats, inds, feat_masks=None):
  feats = feats.view(-1, dim)
  return feats
 
- def _topk(self, scores, K=80):
+ def _topk(self, scores: Tensor, K: int = 80) -> Tuple[Tensor]:
  """Get indexes based on scores.
 
  Args:
@@ -95,7 +101,7 @@ def _topk(self, scores, K=80):
 
  return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
 
- def _transpose_and_gather_feat(self, feat, ind):
+ def _transpose_and_gather_feat(self, feat: Tensor, ind: Tensor) -> Tensor:
  """Given feats and indexes, returns the transposed and gathered feats.
 
  Args:
@@ -115,14 +121,14 @@ def encode(self):
  pass
 
  def decode(self,
- heat,
- rot_sine,
- rot_cosine,
- hei,
- dim,
- vel,
- reg=None,
- task_id=-1):
+ heat: Tensor,
+ rot_sine: Tensor,
+ rot_cosine: Tensor,
+ hei: Tensor,
+ dim: Tensor,
+ vel: Tensor,
+ reg: Tensor = None,
+ task_id: int = -1) -> List[Dict]:
  """Decode bboxes.
 
  Args: