[Feature] Add type hint for code in `models/task_modules` (#2485)

* 2023/04/28 add task_modules type hint

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/coders/groupfree3d_bbox_coder.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* Update mmdet3d/models/task_modules/anchor/anchor_3d_generator.py
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

* 2023/05/19 fix_wrong_hint

* Update centerpoint_bbox_coders.py

* Update iou_neg_piecewise_sampler.py

---------
Co-authored-by: Xiang Xu <xuxiang0103@gmail.com>

mmdet3d/models/task_modules/anchor/anchor_3d_generator.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Tuple, Union
+
 import mmengine
 import torch
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 
@@ -37,13 +40,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 +67,7 @@ class Anchor3DRangeGenerator(object):
         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 +78,21 @@ class Anchor3DRangeGenerator(object):
         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[int]],
+            device: Union[str, torch.device] = 'cuda') -> List[Tensor]:
         """Generate grid anchors in multiple feature levels.
 
         Args:
@@ -112,7 +118,11 @@ class Anchor3DRangeGenerator(object):
             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: int,
+            device: Union[str, torch.device] = 'cuda') -> Tensor:
         """Generate grid anchors of a single level feature map.
 
         This function is usually called by method ``self.grid_anchors``.
@@ -152,13 +162,14 @@ class Anchor3DRangeGenerator(object):
         mr_anchors = torch.cat(mr_anchors, dim=-3)
         return mr_anchors
 
-    def anchors_single_range(self,
-                             feature_size,
-                             anchor_range,
-                             scale=1,
-                             sizes=[[3.9, 1.6, 1.56]],
-                             rotations=[0, 1.5707963],
-                             device='cuda'):
+    def anchors_single_range(
+            self,
+            feature_size: Tuple[int],
+            anchor_range: Union[Tensor, List[float]],
+            scale: int = 1,
+            sizes: Union[List[List[float]], List[float]] = [[3.9, 1.6, 1.56]],
+            rotations: List[float] = [0, 1.5707963],
+            device: Union[str, torch.device] = 'cuda') -> Tensor:
         """Generate anchors in a single range.
 
         Args:
@@ -248,17 +259,18 @@ 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'):
+    def anchors_single_range(
+            self,
+            feature_size: List[int],
+            anchor_range: List[float],
+            scale: int,
+            sizes: Union[List[List[float]], List[float]] = [[3.9, 1.6, 1.56]],
+            rotations: List[float] = [0, 1.5707963],
+            device: Union[str, torch.device] = 'cuda') -> Tensor:
         """Generate anchors in a single range.
 
         Args:
@@ -352,12 +364,15 @@ 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[int]],
+            device: Union[str, torch.device] = 'cuda') -> List[List[Tensor]]:
         """Generate grid anchors in multiple feature levels.
 
         Args:
@@ -379,7 +394,11 @@ class AlignedAnchor3DRangeGeneratorPerCls(AlignedAnchor3DRangeGenerator):
         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[int]],
+            scale: int,
+            device: Union[str, torch.device] = 'cuda') -> List[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

 # 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,7 +42,8 @@ class Max3DIoUAssigner(MaxIoUAssigner):
         iou_calculator (dict): Config of overlaps Calculator.
     """
 
-    def __init__(self,
+    def __init__(
+        self,
         pos_iou_thr: float,
         neg_iou_thr: Union[float, tuple],
         min_pos_iou: float = .0,
@@ -51,7 +52,8 @@ class Max3DIoUAssigner(MaxIoUAssigner):
         ignore_wrt_candidates: bool = True,
         match_low_quality: bool = True,
         gpu_assign_thr: float = -1,
-                 iou_calculator: dict = dict(type='BboxOverlaps2D')):
+        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

 # 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

 # 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 @@ class AnchorFreeBBoxCoder(PartialBinBasedBBoxCoder):
         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 @@ class AnchorFreeBBoxCoder(PartialBinBasedBBoxCoder):
         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

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Optional, 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: Optional[List[float]] = None,
+                 max_num: int = 100,
+                 score_threshold: Optional[float] = None,
+                 code_size: int = 9) -> None:
 
         self.pc_range = pc_range
         self.out_size_factor = out_size_factor
@@ -38,7 +41,10 @@ class CenterPointBBoxCoder(BaseBBoxCoder):
         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: Optional[Tensor] = None) -> Tensor:
         """Given feats and indexes, returns the gathered feats.
 
         Args:
@@ -60,7 +66,7 @@ class CenterPointBBoxCoder(BaseBBoxCoder):
             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 @@ class CenterPointBBoxCoder(BaseBBoxCoder):
 
         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 @@ class CenterPointBBoxCoder(BaseBBoxCoder):
         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: Optional[Tensor] = None,
+               task_id: int = -1) -> List[Dict[str, Tensor]]:
         """Decode bboxes.
 
         Args:

mmdet3d/models/task_modules/coders/delta_xyzwhlr_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 
@@ -13,12 +14,12 @@ class DeltaXYZWLHRBBoxCoder(BaseBBoxCoder):
         code_size (int): The dimension of boxes to be encoded.
     """
 
-    def __init__(self, code_size=7):
+    def __init__(self, code_size: int = 7) -> None:
         super(DeltaXYZWLHRBBoxCoder, self).__init__()
         self.code_size = code_size
 
     @staticmethod
-    def encode(src_boxes, dst_boxes):
+    def encode(src_boxes: Tensor, dst_boxes: Tensor) -> Tensor:
         """Get box regression transformation deltas (dx, dy, dz, dx_size,
         dy_size, dz_size, dr, dv*) that can be used to transform the
         `src_boxes` into the `target_boxes`.
@@ -55,7 +56,7 @@ class DeltaXYZWLHRBBoxCoder(BaseBBoxCoder):
         return torch.cat([xt, yt, zt, wt, lt, ht, rt, *cts], dim=-1)
 
     @staticmethod
-    def decode(anchors, deltas):
+    def decode(anchors: Tensor, deltas: Tensor) -> Tensor:
         """Apply transformation `deltas` (dx, dy, dz, dx_size, dy_size,
         dz_size, dr, dv*) to `boxes`.
 

mmdet3d/models/task_modules/coders/fcos3d_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Optional, Tuple
+
 import numpy as np
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 from mmdet3d.structures.bbox_3d import limit_period
@@ -22,10 +25,10 @@ class FCOS3DBBoxCoder(BaseBBoxCoder):
     """
 
     def __init__(self,
-                 base_depths=None,
-                 base_dims=None,
-                 code_size=7,
-                 norm_on_bbox=True):
+                 base_depths: Optional[Tuple[Tuple[float]]] = None,
+                 base_dims: Optional[Tuple[Tuple[float]]] = None,
+                 code_size: int = 7,
+                 norm_on_bbox: bool = True) -> None:
         super(FCOS3DBBoxCoder, self).__init__()
         self.base_depths = base_depths
         self.base_dims = base_dims
@@ -36,7 +39,12 @@ class FCOS3DBBoxCoder(BaseBBoxCoder):
         # TODO: refactor the encoder in the FCOS3D and PGD head
         pass
 
-    def decode(self, bbox, scale, stride, training, cls_score=None):
+    def decode(self,
+               bbox: Tensor,
+               scale: tuple,
+               stride: int,
+               training: bool,
+               cls_score: Optional[Tensor] = None) -> Tensor:
         """Decode regressed results into 3D predictions.
 
         Note that offsets are not transformed to the projected 3D centers.
@@ -100,7 +108,8 @@ class FCOS3DBBoxCoder(BaseBBoxCoder):
         return bbox
 
     @staticmethod
-    def decode_yaw(bbox, centers2d, dir_cls, dir_offset, cam2img):
+    def decode_yaw(bbox: Tensor, centers2d: Tensor, dir_cls: Tensor,
+                   dir_offset: float, cam2img: Tensor) -> Tensor:
         """Decode yaw angle and change it from local to global.i.
 
         Args:

mmdet3d/models/task_modules/coders/groupfree3d_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List
+
 import numpy as np
 import torch
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures.bbox_3d import BaseInstance3DBoxes
 from .partial_bin_based_bbox_coder import PartialBinBasedBBoxCoder
 
 
@@ -21,11 +25,11 @@ class GroupFree3DBBoxCoder(PartialBinBasedBBoxCoder):
     """
 
     def __init__(self,
-                 num_dir_bins,
-                 num_sizes,
-                 mean_sizes,
-                 with_rot=True,
-                 size_cls_agnostic=True):
+                 num_dir_bins: int,
+                 num_sizes: int,
+                 mean_sizes: List[List[int]],
+                 with_rot: bool = True,
+                 size_cls_agnostic: bool = True) -> None:
         super(GroupFree3DBBoxCoder, self).__init__(
             num_dir_bins=num_dir_bins,
             num_sizes=num_sizes,
@@ -33,7 +37,8 @@ class GroupFree3DBBoxCoder(PartialBinBasedBBoxCoder):
             with_rot=with_rot)
         self.size_cls_agnostic = size_cls_agnostic
 
-    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:
@@ -65,7 +70,7 @@ class GroupFree3DBBoxCoder(PartialBinBasedBBoxCoder):
         return (center_target, size_target, size_class_target, size_res_target,
                 dir_class_target, dir_res_target)
 
-    def decode(self, bbox_out, prefix=''):
+    def decode(self, bbox_out: dict, prefix: str = '') -> Tensor:
         """Decode predicted parts to bbox3d.
 
         Args:
@@ -116,7 +121,11 @@ class GroupFree3DBBoxCoder(PartialBinBasedBBoxCoder):
         bbox3d = torch.cat([center, bbox_size, dir_angle], dim=-1)
         return bbox3d
 
-    def split_pred(self, cls_preds, reg_preds, base_xyz, prefix=''):
+    def split_pred(self,
+                   cls_preds: Tensor,
+                   reg_preds: Tensor,
+                   base_xyz: Tensor,
+                   prefix: str = '') -> Dict[str, Tensor]:
         """Split predicted features to specific parts.
 
         Args:

mmdet3d/models/task_modules/coders/monoflex_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List, Tuple
+
 import numpy as np
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 from torch.nn import functional as F
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures.bbox_3d import BaseInstance3DBoxes
 
 
 @TASK_UTILS.register_module()
@@ -35,19 +39,19 @@ class MonoFlexCoder(BaseBBoxCoder):
     """
 
     def __init__(self,
-                 depth_mode,
-                 base_depth,
-                 depth_range,
-                 combine_depth,
-                 uncertainty_range,
-                 base_dims,
-                 dims_mode,
-                 multibin,
-                 num_dir_bins,
-                 bin_centers,
-                 bin_margin,
-                 code_size,
-                 eps=1e-3):
+                 depth_mode: str,
+                 base_depth: Tuple[float],
+                 depth_range: list,
+                 combine_depth: bool,
+                 uncertainty_range: list,
+                 base_dims: Tuple[Tuple[float]],
+                 dims_mode: str,
+                 multibin: bool,
+                 num_dir_bins: int,
+                 bin_centers: List[float],
+                 bin_margin: float,
+                 code_size: int,
+                 eps: float = 1e-3) -> None:
         super(MonoFlexCoder, self).__init__()
 
         # depth related
@@ -71,7 +75,7 @@ class MonoFlexCoder(BaseBBoxCoder):
         self.bbox_code_size = code_size
         self.eps = eps
 
-    def encode(self, gt_bboxes_3d):
+    def encode(self, gt_bboxes_3d: BaseInstance3DBoxes) -> Tensor:
         """Encode ground truth to prediction targets.
 
         Args:
@@ -105,7 +109,8 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return orientation_target
 
-    def decode(self, bbox, base_centers2d, labels, downsample_ratio, cam2imgs):
+    def decode(self, bbox: Tensor, base_centers2d: Tensor, labels: Tensor,
+               downsample_ratio: int, cam2imgs: Tensor) -> Dict[str, Tensor]:
         """Decode bounding box regression into 3D predictions.
 
         Args:
@@ -211,7 +216,7 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return preds
 
-    def decode_direct_depth(self, depth_offsets):
+    def decode_direct_depth(self, depth_offsets: Tensor) -> Tensor:
         """Transform depth offset to directly regressed depth.
 
         Args:
@@ -239,12 +244,12 @@ class MonoFlexCoder(BaseBBoxCoder):
         return direct_depth
 
     def decode_location(self,
-                        base_centers2d,
-                        offsets2d,
-                        depths,
-                        cam2imgs,
-                        downsample_ratio,
-                        pad_mode='default'):
+                        base_centers2d: Tensor,
+                        offsets2d: Tensor,
+                        depths: Tensor,
+                        cam2imgs: Tensor,
+                        downsample_ratio: Tensor,
+                        pad_mode: str = 'default') -> Tuple[Tensor]:
         """Retrieve object location.
 
         Args:
@@ -283,13 +288,15 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return locations[:, :3]
 
-    def keypoints2depth(self,
-                        keypoints2d,
-                        dimensions,
-                        cam2imgs,
-                        downsample_ratio=4,
-                        group0_index=[(7, 3), (0, 4)],
-                        group1_index=[(2, 6), (1, 5)]):
+    def keypoints2depth(
+            self,
+            keypoints2d: Tensor,
+            dimensions: Tensor,
+            cam2imgs: Tensor,
+            downsample_ratio: int = 4,
+            group0_index: List[Tuple[int]] = [(7, 3), (0, 4)],
+            group1_index: List[Tuple[int]] = [(2, 6),
+                                              (1, 5)]) -> Tuple[Tensor]:
         """Decode depth form three groups of keypoints and geometry projection
         model. 2D keypoints inlucding 8 coreners and top/bottom centers will be
         divided into three groups which will be used to calculate three depths
@@ -383,7 +390,7 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return keypoints_depth
 
-    def decode_dims(self, labels, dims_offset):
+    def decode_dims(self, labels: Tensor, dims_offset: Tensor) -> Tensor:
         """Retrieve object dimensions.
 
         Args:
@@ -411,7 +418,8 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return dimensions
 
-    def decode_orientation(self, ori_vector, locations):
+    def decode_orientation(self, ori_vector: Tensor,
+                           locations: Tensor) -> Tuple[Tensor]:
         """Retrieve object orientation.
 
         Args:
@@ -467,7 +475,8 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return yaws, local_yaws
 
-    def decode_bboxes2d(self, reg_bboxes2d, base_centers2d):
+    def decode_bboxes2d(self, reg_bboxes2d: Tensor,
+                        base_centers2d: Tensor) -> Tensor:
         """Retrieve [x1, y1, x2, y2] format 2D bboxes.
 
         Args:
@@ -492,7 +501,8 @@ class MonoFlexCoder(BaseBBoxCoder):
 
         return bboxes2d
 
-    def combine_depths(self, depth, depth_uncertainty):
+    def combine_depths(self, depth: Tensor,
+                       depth_uncertainty: Tensor) -> Tensor:
         """Combine all the prediced depths with depth uncertainty.
 
         Args:

mmdet3d/models/task_modules/coders/partial_bin_based_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, List
+
 import numpy as np
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures.bbox_3d import BaseInstance3DBoxes
 
 
 @TASK_UTILS.register_module()
@@ -17,7 +21,11 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
         with_rot (bool): Whether the bbox is with rotation.
     """
 
-    def __init__(self, num_dir_bins, num_sizes, mean_sizes, with_rot=True):
+    def __init__(self,
+                 num_dir_bins: int,
+                 num_sizes: int,
+                 mean_sizes: List[List[int]],
+                 with_rot: bool = True):
         super(PartialBinBasedBBoxCoder, self).__init__()
         assert len(mean_sizes) == num_sizes
         self.num_dir_bins = num_dir_bins
@@ -25,7 +33,8 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
         self.mean_sizes = mean_sizes
         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:
@@ -56,7 +65,7 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
         return (center_target, size_class_target, size_res_target,
                 dir_class_target, dir_res_target)
 
-    def decode(self, bbox_out, suffix=''):
+    def decode(self, bbox_out: dict, suffix: str = '') -> Tensor:
         """Decode predicted parts to bbox3d.
 
         Args:
@@ -99,7 +108,8 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
         bbox3d = torch.cat([center, bbox_size, dir_angle], dim=-1)
         return bbox3d
 
-    def decode_corners(self, center, size_res, size_class):
+    def decode_corners(self, center: Tensor, size_res: Tensor,
+                       size_class: Tensor) -> Tensor:
         """Decode center, size residuals and class to corners. Only useful for
         axis-aligned bounding boxes, so angle isn't considered.
 
@@ -137,7 +147,8 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
         corners = torch.cat([corner1, corner2], dim=-1)
         return corners
 
-    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:
@@ -201,7 +212,7 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
 
         return results
 
-    def angle2class(self, angle):
+    def angle2class(self, angle: Tensor) -> tuple:
         """Convert continuous angle to a discrete class and a residual.
 
         Convert continuous angle to a discrete class and a small
@@ -222,7 +233,10 @@ class PartialBinBasedBBoxCoder(BaseBBoxCoder):
             angle_cls * angle_per_class + angle_per_class / 2)
         return angle_cls.long(), angle_res
 
-    def class2angle(self, angle_cls, angle_res, limit_period=True):
+    def class2angle(self,
+                    angle_cls: Tensor,
+                    angle_res: Tensor,
+                    limit_period: bool = True) -> Tensor:
         """Inverse function to angle2class.
 
         Args:

mmdet3d/models/task_modules/coders/pgd_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import Tuple
+
 import numpy as np
 import torch
+from torch import Tensor
 from torch.nn import functional as F
 
 from mmdet3d.registry import TASK_UTILS
@@ -16,13 +19,13 @@ class PGDBBoxCoder(FCOS3DBBoxCoder):
         pass
 
     def decode_2d(self,
-                  bbox,
-                  scale,
-                  stride,
-                  max_regress_range,
-                  training,
-                  pred_keypoints=False,
-                  pred_bbox2d=True):
+                  bbox: Tensor,
+                  scale: tuple,
+                  stride: int,
+                  max_regress_range: int,
+                  training: bool,
+                  pred_keypoints: bool = False,
+                  pred_bbox2d: bool = True) -> Tensor:
         """Decode regressed 2D attributes.
 
         Args:
@@ -70,8 +73,9 @@ class PGDBBoxCoder(FCOS3DBBoxCoder):
                 bbox[:, -4:] = bbox.clone()[:, -4:].exp()
         return bbox
 
-    def decode_prob_depth(self, depth_cls_preds, depth_range, depth_unit,
-                          division, num_depth_cls):
+    def decode_prob_depth(self, depth_cls_preds: Tensor,
+                          depth_range: Tuple[float], depth_unit: int,
+                          division: str, num_depth_cls: int) -> Tensor:
         """Decode probabilistic depth map.
 
         Args:

mmdet3d/models/task_modules/coders/point_xyzwhlr_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Optional
+
 import numpy as np
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures import BaseInstance3DBoxes
 
 
 @TASK_UTILS.register_module()
@@ -18,7 +22,10 @@ class PointXYZWHLRBBoxCoder(BaseBBoxCoder):
             each class. Defaults to None.
     """
 
-    def __init__(self, code_size=7, use_mean_size=True, mean_size=None):
+    def __init__(self,
+                 code_size: int = 7,
+                 use_mean_size: bool = True,
+                 mean_size: List[List[float]] = None):
         super(PointXYZWHLRBBoxCoder, self).__init__()
         self.code_size = code_size
         self.use_mean_size = use_mean_size
@@ -28,7 +35,10 @@ class PointXYZWHLRBBoxCoder(BaseBBoxCoder):
                 f'The min of mean_size should > 0, however currently it is '\
                 f'{self.mean_size.min()}, please check it in your config.'
 
-    def encode(self, gt_bboxes_3d, points, gt_labels_3d=None):
+    def encode(self,
+               gt_bboxes_3d: BaseInstance3DBoxes,
+               points: Tensor,
+               gt_labels_3d: Optional[Tensor] = None) -> Tensor:
         """Encode ground truth to prediction targets.
 
         Args:
@@ -75,7 +85,10 @@ class PointXYZWHLRBBoxCoder(BaseBBoxCoder):
              torch.sin(rg), *cgs],
             dim=-1)
 
-    def decode(self, box_encodings, points, pred_labels_3d=None):
+    def decode(self,
+               box_encodings: Tensor,
+               points: Tensor,
+               pred_labels_3d: Optional[Tensor] = None) -> Tensor:
         """Decode predicted parts and points to bbox3d.
 
         Args:

mmdet3d/models/task_modules/coders/smoke_bbox_coder.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Optional, Tuple
+
 import numpy as np
 import torch
 from mmdet.models.task_modules import BaseBBoxCoder
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
+from mmdet3d.structures import CameraInstance3DBoxes
 
 
 @TASK_UTILS.register_module()
@@ -17,13 +21,16 @@ class SMOKECoder(BaseBBoxCoder):
         code_size (int): The dimension of boxes to be encoded.
     """
 
-    def __init__(self, base_depth, base_dims, code_size):
+    def __init__(self, base_depth: Tuple[float], base_dims: Tuple[float],
+                 code_size: int):
         super(SMOKECoder, self).__init__()
         self.base_depth = base_depth
         self.base_dims = base_dims
         self.bbox_code_size = code_size
 
-    def encode(self, locations, dimensions, orientations, input_metas):
+    def encode(self, locations: Optional[Tensor], dimensions: Tensor,
+               orientations: Tensor,
+               input_metas: List[dict]) -> CameraInstance3DBoxes:
         """Encode CameraInstance3DBoxes by locations, dimensions, orientations.
 
         Args:
@@ -50,12 +57,12 @@ class SMOKECoder(BaseBBoxCoder):
         return batch_bboxes
 
     def decode(self,
-               reg,
-               points,
-               labels,
-               cam2imgs,
-               trans_mats,
-               locations=None):
+               reg: Tensor,
+               points: Tensor,
+               labels: Tensor,
+               cam2imgs: Tensor,
+               trans_mats: Tensor,
+               locations: Optional[Tensor] = None) -> Tuple[Tensor]:
         """Decode regression into locations, dimensions, orientations.
 
         Args:
@@ -104,15 +111,16 @@ class SMOKECoder(BaseBBoxCoder):
 
         return pred_locations, pred_dimensions, pred_orientations
 
-    def _decode_depth(self, depth_offsets):
+    def _decode_depth(self, depth_offsets: Tensor) -> Tensor:
         """Transform depth offset to depth."""
         base_depth = depth_offsets.new_tensor(self.base_depth)
         depths = depth_offsets * base_depth[1] + base_depth[0]
 
         return depths
 
-    def _decode_location(self, points, centers2d_offsets, depths, cam2imgs,
-                         trans_mats):
+    def _decode_location(self, points: Tensor, centers2d_offsets: Tensor,
+                         depths: Tensor, cam2imgs: Tensor,
+                         trans_mats: Tensor) -> Tensor:
         """Retrieve objects location in camera coordinate based on projected
         points.
 
@@ -152,7 +160,7 @@ class SMOKECoder(BaseBBoxCoder):
 
         return locations[:, :3]
 
-    def _decode_dimension(self, labels, dims_offset):
+    def _decode_dimension(self, labels: Tensor, dims_offset: Tensor) -> Tensor:
         """Transform dimension offsets to dimension according to its category.
 
         Args:
@@ -168,7 +176,8 @@ class SMOKECoder(BaseBBoxCoder):
 
         return dimensions
 
-    def _decode_orientation(self, ori_vector, locations):
+    def _decode_orientation(self, ori_vector: Tensor,
+                            locations: Optional[Tensor]) -> Tensor:
         """Retrieve object orientation.
 
         Args:

mmdet3d/models/task_modules/samplers/iou_neg_piecewise_sampler.py

 # Copyright (c) OpenMMLab. All rights reserved.
 import math
+from typing import Optional, Union
 
 import torch
+from mmdet.models.task_modules import AssignResult
+from numpy import ndarray
+from torch import Tensor
 
 from mmdet3d.registry import TASK_UTILS
 from . import RandomSampler, SamplingResult
@@ -29,13 +33,13 @@ class IoUNegPiecewiseSampler(RandomSampler):
     """
 
     def __init__(self,
-                 num,
-                 pos_fraction=None,
-                 neg_piece_fractions=None,
-                 neg_iou_piece_thrs=None,
-                 neg_pos_ub=-1,
-                 add_gt_as_proposals=False,
-                 return_iou=False):
+                 num: int,
+                 pos_fraction: Optional[float] = None,
+                 neg_piece_fractions: Optional[list] = None,
+                 neg_iou_piece_thrs: Optional[list] = None,
+                 neg_pos_ub: float = -1,
+                 add_gt_as_proposals: bool = False,
+                 return_iou: bool = False) -> None:
         super(IoUNegPiecewiseSampler,
               self).__init__(num, pos_fraction, neg_pos_ub,
                              add_gt_as_proposals)
@@ -46,7 +50,8 @@ class IoUNegPiecewiseSampler(RandomSampler):
         self.return_iou = return_iou
         self.neg_piece_num = len(self.neg_piece_fractions)
 
-    def _sample_pos(self, assign_result, num_expected, **kwargs):
+    def _sample_pos(self, assign_result: AssignResult, num_expected: int,
+                    **kwargs) -> Union[Tensor, ndarray]:
         """Randomly sample some positive samples."""
         pos_inds = torch.nonzero(assign_result.gt_inds > 0, as_tuple=False)
         if pos_inds.numel() != 0:
@@ -56,7 +61,8 @@ class IoUNegPiecewiseSampler(RandomSampler):
         else:
             return self.random_choice(pos_inds, num_expected)
 
-    def _sample_neg(self, assign_result, num_expected, **kwargs):
+    def _sample_neg(self, assign_result: AssignResult, num_expected: int,
+                    **kwargs) -> Tensor:
         """Randomly sample some negative samples."""
         neg_inds = torch.nonzero(assign_result.gt_inds == 0, as_tuple=False)
         if neg_inds.numel() != 0:
@@ -127,11 +133,11 @@ class IoUNegPiecewiseSampler(RandomSampler):
             return neg_inds_choice
 
     def sample(self,
-               assign_result,
-               bboxes,
-               gt_bboxes,
-               gt_labels=None,
-               **kwargs):
+               assign_result: AssignResult,
+               bboxes: Tensor,
+               gt_bboxes: Tensor,
+               gt_labels: Optional[Tensor] = None,
+               **kwargs) -> SamplingResult:
         """Sample positive and negative bboxes.
 
         This is a simple implementation of bbox sampling given candidates,

mmdet3d/models/task_modules/samplers/pseudosample.py

@@ -25,7 +25,7 @@ class PseudoSampler(BaseSampler):
         raise NotImplementedError
 
     def sample(self, assign_result: AssignResult, pred_instances: InstanceData,
-               gt_instances: InstanceData, *args, **kwargs):
+               gt_instances: InstanceData, *args, **kwargs) -> SamplingResult:
         """Directly returns the positive and negative indices  of samples.
 
         Args:

mmdet3d/models/task_modules/voxel/voxel_generator.py

 # Copyright (c) OpenMMLab. All rights reserved.
+from typing import List, Tuple, Union
+
 import numba
 import numpy as np
 
@@ -18,10 +20,10 @@ class VoxelGenerator(object):
     """
 
     def __init__(self,
-                 voxel_size,
-                 point_cloud_range,
-                 max_num_points,
-                 max_voxels=20000):
+                 voxel_size: List[float],
+                 point_cloud_range: List[float],
+                 max_num_points: int,
+                 max_voxels: int = 20000):
 
         point_cloud_range = np.array(point_cloud_range, dtype=np.float32)
         # [0, -40, -3, 70.4, 40, 1]
@@ -36,33 +38,33 @@ class VoxelGenerator(object):
         self._max_voxels = max_voxels
         self._grid_size = grid_size
 
-    def generate(self, points):
+    def generate(self, points: np.ndarray) -> Tuple[np.ndarray]:
         """Generate voxels given points."""
         return points_to_voxel(points, self._voxel_size,
                                self._point_cloud_range, self._max_num_points,
                                True, self._max_voxels)
 
     @property
-    def voxel_size(self):
+    def voxel_size(self) -> List[float]:
         """list[float]: Size of a single voxel."""
         return self._voxel_size
 
     @property
-    def max_num_points_per_voxel(self):
+    def max_num_points_per_voxel(self) -> int:
         """int: Maximum number of points per voxel."""
         return self._max_num_points
 
     @property
-    def point_cloud_range(self):
+    def point_cloud_range(self) -> List[float]:
         """list[float]: Range of point cloud."""
         return self._point_cloud_range
 
     @property
-    def grid_size(self):
+    def grid_size(self) -> np.ndarray:
         """np.ndarray: The size of grids."""
         return self._grid_size
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """str: Return a string that describes the module."""
         repr_str = self.__class__.__name__
         indent = ' ' * (len(repr_str) + 1)
@@ -76,12 +78,13 @@ class VoxelGenerator(object):
         return repr_str
 
 
-def points_to_voxel(points,
-                    voxel_size,
-                    coors_range,
-                    max_points=35,
-                    reverse_index=True,
-                    max_voxels=20000):
+def points_to_voxel(points: np.ndarray,
+                    voxel_size: Union[list, tuple, np.ndarray],
+                    coors_range: Union[List[float], List[Tuple[float]],
+                                       List[np.ndarray]],
+                    max_points: int = 35,
+                    reverse_index: bool = True,
+                    max_voxels: int = 20000) -> Tuple[np.ndarray]:
     """convert kitti points(N, >=3) to voxels.
 
     Args:
@@ -138,15 +141,17 @@ def points_to_voxel(points,
 
 
 @numba.jit(nopython=True)
-def _points_to_voxel_reverse_kernel(points,
-                                    voxel_size,
-                                    coors_range,
-                                    num_points_per_voxel,
-                                    coor_to_voxelidx,
-                                    voxels,
-                                    coors,
-                                    max_points=35,
-                                    max_voxels=20000):
+def _points_to_voxel_reverse_kernel(points: np.ndarray,
+                                    voxel_size: Union[list, tuple, np.ndarray],
+                                    coors_range: Union[List[float],
+                                                       List[Tuple[float]],
+                                                       List[np.ndarray]],
+                                    num_points_per_voxel: int,
+                                    coor_to_voxelidx: np.ndarray,
+                                    voxels: np.ndarray,
+                                    coors: np.ndarray,
+                                    max_points: int = 35,
+                                    max_voxels: int = 20000):
     """convert kitti points(N, >=3) to voxels.
 
     Args:
@@ -212,15 +217,16 @@ def _points_to_voxel_reverse_kernel(points,
 
 
 @numba.jit(nopython=True)
-def _points_to_voxel_kernel(points,
-                            voxel_size,
-                            coors_range,
-                            num_points_per_voxel,
-                            coor_to_voxelidx,
-                            voxels,
-                            coors,
-                            max_points=35,
-                            max_voxels=20000):
+def _points_to_voxel_kernel(points: np.ndarray,
+                            voxel_size: Union[list, tuple, np.ndarray],
+                            coors_range: Union[List[float], List[Tuple[float]],
+                                               List[np.ndarray]],
+                            num_points_per_voxel: int,
+                            coor_to_voxelidx: np.ndarray,
+                            voxels: np.ndarray,
+                            coors: np.ndarray,
+                            max_points: int = 35,
+                            max_voxels: int = 200000):
     """convert kitti points(N, >=3) to voxels.
 
     Args:
@@ -230,7 +236,7 @@ def _points_to_voxel_kernel(points,
         coors_range (list[float | tuple[float] | ndarray]): Range of voxels.
             format: xyzxyz, minmax
         num_points_per_voxel (int): Number of points per voxel.
-        coor_to_voxel_idx (np.ndarray): A voxel grid of shape (D, H, W),
+        coor_to_voxelidx (np.ndarray): A voxel grid of shape (D, H, W),
             which has the same shape as the complete voxel map. It indicates
             the index of each corresponding voxel.
         voxels (np.ndarray): Created empty voxels.