[Refactor] CSL (#535)

* add UT

* fix lint

* update

* Update test_oriented_reppoints.py

* fix UT

* init

* Update angle_branch_retina_head.py

* add ut

* Update test_angle_coder.py

* add UT

* Update test_angle_branch_retina_head.py

* fix

* fix

Author: zytx121

configs/csl/rotated_retinanet_obb_csl_gaussian_r50_fpn_fp16_1x_dota_le90.py

@@ -4,19 +4,15 @@
 angle_version = 'le90'
 model = dict(
     bbox_head=dict(
-        type='CSLRRetinaHead',
+        type='AngleBranchRetinaHead',
         angle_coder=dict(
             type='CSLCoder',
             angle_version=angle_version,
             omega=4,
             window='gaussian',
             radius=3),
-        loss_cls=dict(
-            type='FocalLoss',
-            use_sigmoid=True,
+        loss_angle=dict(
+            type='mmdet.SmoothFocalLoss',
             gamma=2.0,
             alpha=0.25,
-            loss_weight=1.0),
-        loss_bbox=dict(type='L1Loss', loss_weight=1.0),
-        loss_angle=dict(
-            type='SmoothFocalLoss', gamma=2.0, alpha=0.25, loss_weight=0.8)))
+            loss_weight=0.8)))

configs/oriented_rcnn/oriented_rcnn_r50_fpn_fp16_1x_dota_le90.py

@@ -1,3 +1,3 @@
 _base_ = './oriented_rcnn_r50_fpn_1x_dota_le90.py'
 
-optim_wrapper = dict(type='AmpOptimWrapper')
+optim_wrapper = dict(type='AmpOptimWrapper', loss_scale='dynamic')

configs/rotated_retinanet/rotated_retinanet_obb_r50_fpn_fp16_1x_dota_le90.py

@@ -1,3 +1,3 @@
 _base_ = ['./rotated_retinanet_obb_r50_fpn_1x_dota_le90.py']
 
-fp16 = dict(loss_scale='dynamic')
+optim_wrapper = dict(type='AmpOptimWrapper', loss_scale='dynamic')

mmrotate/core/bbox/coder/angle_coder.py

@@ -3,6 +3,7 @@
 
 import torch
 from mmdet.models.task_modules.coders.base_bbox_coder import BaseBBoxCoder
+from torch import Tensor
 
 from mmrotate.registry import TASK_UTILS
 
@@ -35,44 +36,44 @@ def __init__(self, angle_version, omega=1, window='gaussian', radius=6):
         self.omega = omega
         self.window = window
         self.radius = radius
-        self.coding_len = int(self.angle_range // omega)
+        self.encoded_size = int(self.angle_range // omega)
 
-    def encode(self, angle_targets):
+    def encode(self, angle_targets: Tensor) -> Tensor:
         """Circular Smooth Label Encoder.
 
         Args:
             angle_targets (Tensor): Angle offset for each scale level
                 Has shape (num_anchors * H * W, 1)
 
         Returns:
-            list[Tensor]: The csl encoding of angle offset for each
-                scale level. Has shape (num_anchors * H * W, coding_len)
+            Tensor: The csl encoding of angle offset for each scale
+            level. Has shape (num_anchors * H * W, encoded_size)
         """
 
         # radius to degree
         angle_targets_deg = angle_targets * (180 / math.pi)
         # empty label
         smooth_label = torch.zeros_like(angle_targets).repeat(
-            1, self.coding_len)
+            1, self.encoded_size)
         angle_targets_deg = (angle_targets_deg +
                              self.angle_offset) / self.omega
         # Float to Int
         angle_targets_long = angle_targets_deg.long()
 
         if self.window == 'pulse':
-            radius_range = angle_targets_long % self.coding_len
+            radius_range = angle_targets_long % self.encoded_size
             smooth_value = 1.0
         elif self.window == 'rect':
             base_radius_range = torch.arange(
                 -self.radius, self.radius, device=angle_targets_long.device)
             radius_range = (base_radius_range +
-                            angle_targets_long) % self.coding_len
+                            angle_targets_long) % self.encoded_size
             smooth_value = 1.0
         elif self.window == 'triangle':
             base_radius_range = torch.arange(
                 -self.radius, self.radius, device=angle_targets_long.device)
             radius_range = (base_radius_range +
-                            angle_targets_long) % self.coding_len
+                            angle_targets_long) % self.encoded_size
             smooth_value = 1.0 - torch.abs(
                 (1 / self.radius) * base_radius_range)
 
@@ -83,7 +84,7 @@ def encode(self, angle_targets):
                 device=angle_targets_long.device)
 
             radius_range = (base_radius_range +
-                            angle_targets_long) % self.coding_len
+                            angle_targets_long) % self.encoded_size
             smooth_value = torch.exp(-torch.pow(base_radius_range, 2) /
                                      (2 * self.radius**2))
 
@@ -96,19 +97,24 @@ def encode(self, angle_targets):
 
         return smooth_label.scatter(1, radius_range, smooth_value)
 
-    def decode(self, angle_preds):
+    def decode(self, angle_preds: Tensor, keepdim: bool = False) -> Tensor:
         """Circular Smooth Label Decoder.
 
         Args:
-            angle_preds (Tensor): The csl encoding of angle offset
-                for each scale level.
-                Has shape (num_anchors * H * W, coding_len)
+            angle_preds (Tensor): The csl encoding of angle offset for each
+                scale level. Has shape (num_anchors * H * W, encoded_size) or
+                (B, num_anchors * H * W, encoded_size)
+            keepdim (bool): Whether the output tensor has dim retained or not.
+
 
         Returns:
-            list[Tensor]: Angle offset for each scale level.
-                Has shape (num_anchors * H * W, 1)
+            Tensor: Angle offset for each scale level. When keepdim is true,
+            return (num_anchors * H * W, 1) or (B, num_anchors * H * W, 1),
+            otherwise (num_anchors * H * W,) or (B, num_anchors * H * W)
         """
-        angle_cls_inds = torch.argmax(angle_preds, dim=1)
+        if angle_preds.shape[0] == 0:
+            return angle_preds.new_zeros((0))
+        angle_cls_inds = torch.argmax(angle_preds, dim=-1, keepdim=keepdim)
         angle_pred = ((angle_cls_inds + 0.5) *
                       self.omega) % self.angle_range - self.angle_offset
         return angle_pred * (math.pi / 180)

mmrotate/models/dense_heads/__init__.py

@@ -1,7 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from .angle_branch_retina_head import AngleBranchRetinaHead
 from .cfa_head import CFAHead
 from .csl_rotated_fcos_head import CSLRFCOSHead
-from .csl_rotated_retina_head import CSLRRetinaHead
 from .kfiou_odm_refine_head import KFIoUODMRefineHead
 from .kfiou_rotate_retina_head import KFIoURRetinaHead
 from .kfiou_rotate_retina_refine_head import KFIoURRetinaRefineHead
@@ -23,7 +23,7 @@
     'RotatedRetinaHead', 'RotatedRPNHead', 'OrientedRPNHead',
     'RotatedRetinaRefineHead', 'ODMRefineHead', 'KFIoURRetinaHead',
     'KFIoURRetinaRefineHead', 'KFIoUODMRefineHead', 'RotatedRepPointsHead',
-    'SAMRepPointsHead', 'CSLRRetinaHead', 'RotatedATSSHead',
+    'SAMRepPointsHead', 'AngleBranchRetinaHead', 'RotatedATSSHead',
     'RotatedAnchorFreeHead', 'RotatedFCOSHead', 'CSLRFCOSHead',
     'OrientedRepPointsHead', 'R3Head', 'R3RefineHead', 'S2AHead',
     'S2ARefineHead', 'CFAHead'