[Fix] Fix mmdet ort static shape bug.

mmdeploy/codebase/mmdet/core/bbox/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .delta_xywh_bbox_coder import *  # noqa: F401,F403
+from .distance_point_bbox_coder import *  # noqa: F401,F403
 from .tblr_bbox_coder import *  # noqa: F401,F403
 from .transforms import *  # noqa: F401,F403

mmdeploy/codebase/mmdet/core/bbox/distance_point_bbox_coder.py

@@ -0,0 +1,40 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import mmdet.core.bbox.transforms
+
+from mmdeploy.core import FUNCTION_REWRITER
+
+
+@FUNCTION_REWRITER.register_rewriter(
+    func_name='mmdet.core.bbox.coder.DistancePointBBoxCoder.decode',
+    backend='default')
+def decode__default(ctx, self, points, pred_bboxes, max_shape=None):
+    """Rewrite `mmdet.core.bbox.coder.DistancePointBBoxCoder.decode`
+
+    Decode distance prediction to bounding box.
+
+    Args:
+        ctx (ContextCaller): The context with additional information.
+        self (DistancePointBBoxCoder): The instance of the class
+            DistancePointBBoxCoder.
+        points (Tensor): Shape (B, N, 2) or (N, 2).
+        pred_bboxes (Tensor): Distance from the given point to 4
+            boundaries (left, top, right, bottom). Shape (B, N, 4)
+            or (N, 4)
+        max_shape (Sequence[int] or torch.Tensor or Sequence[
+            Sequence[int]],optional): Maximum bounds for boxes, specifies
+            (H, W, C) or (H, W). If priors shape is (B, N, 4), then
+            the max_shape should be a Sequence[Sequence[int]],
+            and the length of max_shape should also be B.
+            Default None.
+    Returns:
+        Tensor: Boxes with shape (N, 4) or (B, N, 4)
+    """
+    assert points.size(0) == pred_bboxes.size(0)
+    assert points.size(-1) == 2
+    assert pred_bboxes.size(-1) == 4
+    if self.clip_border is False:
+        max_shape = None
+    # Rewrite add mmdet.core.bbox.transforms to find correct
+    # rewriter, or you will not find correct rewriter.
+    return mmdet.core.bbox.transforms.distance2bbox(points, pred_bboxes,
+                                                    max_shape)

mmdeploy/codebase/mmdet/core/bbox/transforms.py

@@ -2,14 +2,19 @@
 import torch
 
 from mmdeploy.codebase.mmdet.deploy import clip_bboxes
+from mmdeploy.core import FUNCTION_REWRITER
 
 
-def distance2bbox(points, distance, max_shape=None):
+@FUNCTION_REWRITER.register_rewriter(
+    func_name='mmdet.core.bbox.transforms.distance2bbox'  # noqa
+)
+def distance2bbox__default(ctx, points, distance, max_shape=None):
     """Rewrite `mmdet.core.bbox.transforms.distance2bbox`
 
     Decode distance prediction to bounding box.
 
     Args:
+        ctx (ContextCaller): The context with additional information.
         points (Tensor): Shape (B, N, 2) or (N, 2).
         distance (Tensor): Distance from the given point to 4
             boundaries (left, top, right, bottom). Shape (B, N, 4) or (N, 4)

mmdeploy/codebase/mmdet/models/dense_heads/base_dense_head.py

@@ -81,7 +81,6 @@ def base_dense_head__get_bbox(ctx,
         mlvl_score_factors = []
     assert img_metas is not None
     img_shape = img_metas[0]['img_shape']
-
     assert len(cls_scores) == len(bbox_preds) == len(mlvl_priors)
     batch_size = cls_scores[0].shape[0]
     cfg = self.test_cfg

tests/test_codebase/test_mmdet/test_mmdet_core.py

@@ -149,12 +149,41 @@ def tblr2bboxes(*args, **kwargs):
         assert rewrite_outputs is not None
 
 
-def test_distance2bbox():
-    from mmdeploy.codebase.mmdet.core import distance2bbox
+@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
+def test_distance2bbox(backend_type: Backend):
+    check_backend(backend_type)
+    deploy_cfg = mmcv.Config(
+        dict(
+            onnx_config=dict(output_names=None, input_shape=None),
+            backend_config=dict(type=backend_type.value, model_inputs=None),
+            codebase_config=dict(type='mmdet', task='ObjectDetection')))
+
+    # wrap function to enable rewrite
+    def distance2bbox(*args, **kwargs):
+        import mmdet.core.bbox.transforms
+        return mmdet.core.bbox.transforms.distance2bbox(*args, **kwargs)
+
     points = torch.rand(3, 2)
     distance = torch.rand(3, 4)
-    bbox = distance2bbox(points, distance)
-    assert bbox.shape == torch.Size([3, 4])
+    original_outputs = distance2bbox(points, distance)
+
+    # wrap function to nn.Module, enable torch.onnx.export
+    wrapped_func = WrapFunction(distance2bbox)
+    rewrite_outputs, is_backend_output = get_rewrite_outputs(
+        wrapped_func,
+        model_inputs={
+            'points': points,
+            'distance': distance
+        },
+        deploy_cfg=deploy_cfg)
+
+    if is_backend_output:
+        model_output = original_outputs.squeeze().cpu().numpy()
+        rewrite_output = rewrite_outputs[0].squeeze()
+        assert np.allclose(
+            model_output, rewrite_output, rtol=1e-03, atol=1e-05)
+    else:
+        assert rewrite_outputs is not None
 
 
 @backend_checker(Backend.ONNXRUNTIME)