feat: add RetinaNet support

tests/test_models/test_forward.py

@@ -631,6 +631,8 @@ def test_inference_detector():
 
     model = build_detector(ConfigDict(model_dict))
     config = _get_config_module("retinanet/retinanet_r50_fpn_1x_coco.py")
+    # Add mock test_dataloader to avoid AttributeError in inference_detector
+    config.test_dataloader = ConfigDict(dataset=ConfigDict(pipeline=config.test_pipeline))
     model.cfg = config
     # test single image
     result = inference_detector(model, img1)

tests/test_models/test_retinanet.py

@@ -0,0 +1,106 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+import torch.nn as nn
+from visdet.structures import DetDataSample, InstanceData
+from visdet.registry import MODELS
+
+
+def test_retinanet_forward_backward():
+    model_cfg = dict(
+        type="RetinaNet",
+        backbone=dict(
+            type="ResNet",
+            depth=18,
+            num_stages=4,
+            out_indices=(0, 1, 2, 3),
+            frozen_stages=1,
+            norm_cfg=dict(type="BN", requires_grad=True),
+            norm_eval=True,
+            style="pytorch",
+        ),
+        neck=dict(
+            type="FPN",
+            in_channels=[64, 128, 256, 512],
+            out_channels=256,
+            start_level=1,
+            add_extra_convs="on_input",
+            num_outs=5,
+        ),
+        bbox_head=dict(
+            type="RetinaHead",
+            num_classes=80,
+            in_channels=256,
+            stacked_convs=4,
+            feat_channels=256,
+            anchor_generator=dict(
+                type="AnchorGenerator",
+                octave_base_scale=4,
+                scales_per_octave=3,
+                ratios=[0.5, 1.0, 2.0],
+                strides=[8, 16, 32, 64, 128],
+            ),
+            bbox_coder=dict(
+                type="DeltaXYWHBBoxCoder", target_means=[0.0, 0.0, 0.0, 0.0], target_stds=[1.0, 1.0, 1.0, 1.0]
+            ),
+            loss_cls=dict(type="FocalLoss", use_sigmoid=True, gamma=2.0, alpha=0.25, loss_weight=1.0),
+            loss_bbox=dict(type="L1Loss", loss_weight=1.0),
+        ),
+        train_cfg=dict(
+            assigner=dict(type="MaxIoUAssigner", pos_iou_thr=0.5, neg_iou_thr=0.4, min_pos_iou=0, ignore_iof_thr=-1),
+            allowed_border=-1,
+            pos_weight=-1,
+            debug=False,
+        ),
+        test_cfg=dict(
+            nms_pre=1000, min_bbox_size=0, score_thr=0.05, nms=dict(type="nms", iou_threshold=0.5), max_per_img=100
+        ),
+    )
+
+    detector = MODELS.build(model_cfg)
+
+    # Test forward train
+    imgs = torch.randn(2, 3, 224, 224, requires_grad=True)
+    data_samples = []
+    for i in range(2):
+        data_sample = DetDataSample()
+        data_sample.set_metainfo(
+            dict(img_shape=(224, 224), ori_shape=(224, 224), pad_shape=(224, 224), scale_factor=(1.0, 1.0))
+        )
+        gt_instances = InstanceData()
+        gt_instances.bboxes = torch.tensor([[10, 10, 50, 50]], dtype=torch.float32)
+        gt_instances.labels = torch.tensor([1], dtype=torch.long)
+        data_sample.gt_instances = gt_instances
+        data_samples.append(data_sample)
+
+    # Mode loss
+    losses = detector(imgs, data_samples, mode="loss")
+    assert isinstance(losses, dict)
+    assert "loss_cls" in losses
+    assert "loss_bbox" in losses
+
+    # Backward pass
+    total_loss = 0
+    for loss_value in losses.values():
+        if isinstance(loss_value, torch.Tensor):
+            total_loss += loss_value.sum()
+        elif isinstance(loss_value, list):
+            total_loss += sum(l.sum() for l in loss_value)
+
+    total_loss.backward()
+    assert imgs.grad is not None
+    print("Backward pass passed!")
+
+    # Test forward predict
+    detector.eval()
+    with torch.no_grad():
+        predictions = detector(imgs, data_samples, mode="predict")
+        assert len(predictions) == 2
+        assert isinstance(predictions[0], DetDataSample)
+        assert hasattr(predictions[0], "pred_instances")
+    print("Forward predict passed!")
+
+    print("RetinaNet forward and backward tests passed!")
+
+
+if __name__ == "__main__":
+    test_retinanet_forward_backward()

visdet/models/dense_heads/__init__.py

@@ -2,5 +2,6 @@
 from visdet.models.dense_heads.base_dense_head import BaseDenseHead
 from visdet.models.dense_heads.anchor_head import AnchorHead
 from visdet.models.dense_heads.rpn_head import RPNHead
+from visdet.models.dense_heads.retina_head import RetinaHead
 
-__all__ = ["BaseDenseHead", "AnchorHead", "RPNHead"]
+__all__ = ["BaseDenseHead", "AnchorHead", "RPNHead", "RetinaHead"]

visdet/models/dense_heads/base_dense_head.py

@@ -466,18 +466,19 @@ def _bbox_post_process(
         # filter small size bboxes
         if cfg.get("min_bbox_size", -1) >= 0:
             w, h = get_box_wh(results.bboxes)
-            valid_mask = (w > cfg.min_bbox_size) & (h > cfg.min_bbox_size)
+            min_bbox_size = cfg.get("min_bbox_size")
+            valid_mask = (w > min_bbox_size) & (h > min_bbox_size)
             if not valid_mask.all():
                 results = results[valid_mask]
 
         # TODO: deal with `with_nms` and `nms_cfg=None` in test_cfg
         if with_nms and results.bboxes.numel() > 0:
             bboxes = get_box_tensor(results.bboxes)
-            det_bboxes, keep_idxs = batched_nms(bboxes, results.scores, results.labels, cfg.nms)
+            det_bboxes, keep_idxs = batched_nms(bboxes, results.scores, results.labels, cfg.get("nms"))
             results = results[keep_idxs]
             # some nms would reweight the score, such as softnms
             results.scores = det_bboxes[:, -1]
-            results = results[: cfg.max_per_img]
+            results = results[: cfg.get("max_per_img")]
 
         return results
 

visdet/models/dense_heads/retina_head.py

@@ -0,0 +1,101 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch.nn as nn
+from visdet.cv.cnn import ConvModule
+
+from visdet.registry import MODELS
+from .anchor_head import AnchorHead
+
+
+@MODELS.register_module()
+class RetinaHead(AnchorHead):
+    """An anchor-based head used in `RetinaNet
+    <https://arxiv.org/abs/1708.02002>`_.
+
+    The head contains two separate stacks of convolution layers, one for
+    classification and one for localization. Each stack has four conv layers
+    with a 3x3 kernel and 256 channels.
+
+    Args:
+        stacked_convs (int): Number of stacking conv layers of the head.
+            Defaults to 4.
+        conv_cfg (dict): Config dict for convolution layer.
+            Defaults to None.
+        norm_cfg (dict): Config dict for normalization layer.
+            Defaults to None.
+        init_cfg (dict or list[dict], optional): Initialization config dict.
+    """
+
+    def __init__(
+        self,
+        num_classes,
+        in_channels,
+        stacked_convs=4,
+        conv_cfg=None,
+        norm_cfg=None,
+        init_cfg=dict(
+            type="Normal",
+            layer="Conv2d",
+            std=0.01,
+            override=dict(type="Normal", name="retina_cls", std=0.01, bias_prob=0.01),
+        ),
+        **kwargs,
+    ):
+        self.stacked_convs = stacked_convs
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        super().__init__(num_classes, in_channels, init_cfg=init_cfg, **kwargs)
+
+    def _init_layers(self):
+        """Initialize layers of the head."""
+        self.relu = nn.ReLU(inplace=True)
+        self.cls_convs = nn.ModuleList()
+        self.reg_convs = nn.ModuleList()
+        for i in range(self.stacked_convs):
+            chn = self.in_channels if i == 0 else self.feat_channels
+            self.cls_convs.append(
+                ConvModule(
+                    chn,
+                    self.feat_channels,
+                    3,
+                    stride=1,
+                    padding=1,
+                    conv_cfg=self.conv_cfg,
+                    norm_cfg=self.norm_cfg,
+                )
+            )
+            self.reg_convs.append(
+                ConvModule(
+                    chn,
+                    self.feat_channels,
+                    3,
+                    stride=1,
+                    padding=1,
+                    conv_cfg=self.conv_cfg,
+                    norm_cfg=self.norm_cfg,
+                )
+            )
+        self.retina_cls = nn.Conv2d(self.feat_channels, self.num_base_priors * self.cls_out_channels, 3, padding=1)
+        self.retina_reg = nn.Conv2d(self.feat_channels, self.num_base_priors * 4, 3, padding=1)
+
+    def forward_single(self, x):
+        """Forward feature of a single scale level.
+
+        Args:
+            x (Tensor): Features of a single scale level.
+
+        Returns:
+            tuple:
+                cls_score (Tensor): Cls scores for a single scale level
+                    the channels number is num_base_priors * num_classes.
+                bbox_pred (Tensor): Box energies / deltas for a single scale
+                    level, the channels number is num_base_priors * 4.
+        """
+        cls_feat = x
+        reg_feat = x
+        for cls_conv in self.cls_convs:
+            cls_feat = cls_conv(cls_feat)
+        for reg_conv in self.reg_convs:
+            reg_feat = reg_conv(reg_feat)
+        cls_score = self.retina_cls(cls_feat)
+        bbox_pred = self.retina_reg(reg_feat)
+        return cls_score, bbox_pred

visdet/models/detectors/__init__.py

@@ -4,6 +4,8 @@
 from visdet.models.detectors.base import BaseDetector
 from visdet.models.detectors.cascade_rcnn import CascadeRCNN
 from visdet.models.detectors.mask_rcnn import MaskRCNN
+from visdet.models.detectors.retinanet import RetinaNet
+from visdet.models.detectors.single_stage import SingleStageDetector
 from visdet.models.detectors.two_stage import TwoStageDetector
 
-__all__ = ["BaseDetector", "TwoStageDetector", "MaskRCNN", "CascadeRCNN"]
+__all__ = ["BaseDetector", "SingleStageDetector", "TwoStageDetector", "RetinaNet", "MaskRCNN", "CascadeRCNN"]

visdet/models/detectors/retinanet.py

@@ -0,0 +1,28 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from visdet.registry import MODELS
+from .single_stage import SingleStageDetector
+
+
+@MODELS.register_module()
+class RetinaNet(SingleStageDetector):
+    """Implementation of `RetinaNet <https://arxiv.org/abs/1708.02002>`_"""
+
+    def __init__(
+        self,
+        backbone,
+        neck,
+        bbox_head,
+        train_cfg=None,
+        test_cfg=None,
+        data_preprocessor=None,
+        init_cfg=None,
+    ):
+        super().__init__(
+            backbone=backbone,
+            neck=neck,
+            bbox_head=bbox_head,
+            train_cfg=train_cfg,
+            test_cfg=test_cfg,
+            data_preprocessor=data_preprocessor,
+            init_cfg=init_cfg,
+        )