From bf36950f0ee583c11162e4b83b8a23d6e926ec13 Mon Sep 17 00:00:00 2001 From: "q.yao" Date: Fri, 24 Feb 2023 10:50:03 +0800 Subject: [PATCH] [Fix] fix centernet (#1768) * fix centernet * update sdk transform --- .../transform/default_format_bundle.cpp | 4 +++- csrc/mmdeploy/preprocess/transform/pad.cpp | 16 +++++++++++++++ mmdeploy/codebase/mmdet/models/__init__.py | 1 + .../codebase/mmdet/models/utils/__init__.py | 2 ++ .../mmdet/models/utils/gaussian_target.py | 20 +++++++++++++++++++ .../test_mmdet/test_mmdet_utils.py | 14 +++++++++++++ 6 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 mmdeploy/codebase/mmdet/models/utils/__init__.py create mode 100644 mmdeploy/codebase/mmdet/models/utils/gaussian_target.py diff --git a/csrc/mmdeploy/preprocess/transform/default_format_bundle.cpp b/csrc/mmdeploy/preprocess/transform/default_format_bundle.cpp index 07aa85d87..3cca32b9d 100644 --- a/csrc/mmdeploy/preprocess/transform/default_format_bundle.cpp +++ b/csrc/mmdeploy/preprocess/transform/default_format_bundle.cpp @@ -36,7 +36,9 @@ class DefaultFormatBundle : public Transform { } } if (!data.contains("scale_factor")) { - data["scale_factor"].push_back(1.0); + for (int i = 0; i < 4; ++i) { + data["scale_factor"].push_back(1.0); + } } if (!data.contains("img_norm_cfg")) { int channel = tensor.shape()[3]; diff --git a/csrc/mmdeploy/preprocess/transform/pad.cpp b/csrc/mmdeploy/preprocess/transform/pad.cpp index fc42f9c64..c02542fae 100644 --- a/csrc/mmdeploy/preprocess/transform/pad.cpp +++ b/csrc/mmdeploy/preprocess/transform/pad.cpp @@ -41,6 +41,10 @@ class Pad : public Transform { } else { pad_val_ = 0.0f; } + + logical_or_val_ = args.value("logical_or_val", 0); + add_pix_val_ = args.value("add_pix_val", 0); + pad_to_square_ = args.value("pad_to_square", false); padding_mode_ = args.value("padding_mode", std::string("constant")); orientation_agnostic_ = args.value("orientation_agnostic", false); @@ -89,6 +93,16 @@ class Pad : public Transform { data["pad_size_divisor"] = size_divisor_; data["pad_fixed_size"].push_back(pad_h); data["pad_fixed_size"].push_back(pad_w); + } else if (logical_or_val_ > 0) { + int pad_h = (height | logical_or_val_) + add_pix_val_; + int pad_w = (width | logical_or_val_) + add_pix_val_; + int offset_h = pad_h / 2 - height / 2; + int offset_w = pad_w / 2 - width / 2; + padding = {offset_w, offset_h, pad_w - width - offset_w, pad_h - height - offset_h}; + data["border"].push_back(offset_h); + data["border"].push_back(offset_w); + data["border"].push_back(offset_h + height); + data["border"].push_back(offset_w + width); } else { output_tensor = tensor; data["pad_fixed_size"].push_back(height); @@ -124,6 +138,8 @@ class Pad : public Transform { operation::Managed pad_; std::array size_; int size_divisor_; + int logical_or_val_; + int add_pix_val_; float pad_val_; bool pad_to_square_; bool orientation_agnostic_; diff --git a/mmdeploy/codebase/mmdet/models/__init__.py b/mmdeploy/codebase/mmdet/models/__init__.py index 38b7e336d..ee7b60a4e 100644 --- a/mmdeploy/codebase/mmdet/models/__init__.py +++ b/mmdeploy/codebase/mmdet/models/__init__.py @@ -7,3 +7,4 @@ from . import necks # noqa: F401,F403 from . import roi_heads # noqa: F401,F403 from . import task_modules # noqa: F401,F403 from . import transformer # noqa: F401,F403 +from . import utils # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmdet/models/utils/__init__.py b/mmdeploy/codebase/mmdet/models/utils/__init__.py new file mode 100644 index 000000000..a5033b73b --- /dev/null +++ b/mmdeploy/codebase/mmdet/models/utils/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from . import gaussian_target # noqa: F401,F403 diff --git a/mmdeploy/codebase/mmdet/models/utils/gaussian_target.py b/mmdeploy/codebase/mmdet/models/utils/gaussian_target.py new file mode 100644 index 000000000..6e39ce99b --- /dev/null +++ b/mmdeploy/codebase/mmdet/models/utils/gaussian_target.py @@ -0,0 +1,20 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import torch + +from mmdeploy.core import FUNCTION_REWRITER + + +@FUNCTION_REWRITER.register_rewriter( + func_name='mmdet.models.utils.gaussian_target.get_topk_from_heatmap') +def get_topk_from_heatmap__default(scores, k=20): + """Get top k positions from heatmap. + + Replace view(batch, -1) with flatten + """ + height, width = scores.size()[2:] + topk_scores, topk_inds = torch.topk(scores.flatten(1), k) + topk_clses = topk_inds // (height * width) + topk_inds = topk_inds % (height * width) + topk_ys = topk_inds // width + topk_xs = (topk_inds % width).int().float() + return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs diff --git a/tests/test_codebase/test_mmdet/test_mmdet_utils.py b/tests/test_codebase/test_mmdet/test_mmdet_utils.py index 2045ad93f..acb9947eb 100644 --- a/tests/test_codebase/test_mmdet/test_mmdet_utils.py +++ b/tests/test_codebase/test_mmdet/test_mmdet_utils.py @@ -68,3 +68,17 @@ def test_get_mmdet_params(): pre_top_k=-1, keep_top_k=100, background_label_id=-1) + + +def test_get_topk_from_heatmap(): + from mmdet.models.utils.gaussian_target import get_topk_from_heatmap + + from mmdeploy.codebase.mmdet.models.utils.gaussian_target import \ + get_topk_from_heatmap__default + scores = torch.rand(1, 2, 4, 4) + + gts = get_topk_from_heatmap(scores, k=20) + outs = get_topk_from_heatmap__default(scores, k=20) + + for gt, out in zip(gts, outs): + torch.testing.assert_allclose(gt, out)