Skip to content

[Fix] Fix transposed prior grid in RTMOHead.switch_to_deploy for non-square input - #3267

Open
bolun365 wants to merge 1 commit into
open-mmlab:mainfrom
bolun365:main
Open

[Fix] Fix transposed prior grid in RTMOHead.switch_to_deploy for non-square input#3267
bolun365 wants to merge 1 commit into
open-mmlab:mainfrom
bolun365:main

Conversation

@bolun365

Copy link
Copy Markdown

Motivation

When deploying RTMO with a non-square input size (e.g. 256x160, W×H),
the exported model (ONNX / MNN / TensorRT) produces wrong detections: a single
person yields many scattered high-confidence "phantom" boxes that NMS cannot
merge, even though the PyTorch model runs correctly.

The root cause is in RTMOHead.switch_to_deploy. It builds the dummy feature
maps used to precompute flatten_priors as:

torch.rand(1, 1, input_size[0] // s, input_size[1] // s)  # (W//s, H//s)

input_size is in (W, H) order, but a real feature map has shape
(H//s, W//s). For square inputs (e.g. 256x256, the default RTMO configs)
the two are equal so the bug is invisible. For non-square inputs the prior grid
is transposed, so boxes are decoded at the wrong locations during export.
(The runtime predict path is unaffected because it derives feature-map sizes
from the actual cls_score.shape[2:].)

Modification

In mmpose/models/heads/hybrid_heads/rtmo_head.py, RTMOHead.switch_to_deploy,
swap the two index orders so the dummy feature map matches the real
(H//s, W//s) layout:

torch.rand(1, 1, input_size[1] // s, input_size[0] // s)  # (H//s, W//s)

BC-breaking (Optional)

No. The change is a no-op for square input_size (all existing official RTMO
configs use square inputs), and only corrects the behavior for non-square
inputs, which was previously broken.


Use cases (Optional)

Deploying RTMO with a non-square (landscape/portrait) input_size, e.g.
input_size=(256, 160), now exports an ONNX/MNN/TensorRT model whose decoded
boxes match the PyTorch results.


## Checklist

**Before PR**:

- [ ] I have read and followed the workflow indicated in the [CONTRIBUTING.md](https://github.com/open-mmlab/mmpose/blob/master/.github/CONTRIBUTING.md) to create this PR.
- [ ] Pre-commit or linting tools indicated in [CONTRIBUTING.md](https://github.com/open-mmlab/mmpose/blob/master/.github/CONTRIBUTING.md) are used to fix the potential lint issues.
- [ ] Bug fixes are covered by unit tests, the case that causes the bug should be added in the unit tests.
- [ ] New functionalities are covered by complete unit tests. If not, please add more unit tests to ensure correctness.
- [ ] The documentation has been modified accordingly, including docstring or example tutorials.

**After PR**:

- [ ] CLA has been signed and all committers have signed the CLA in this PR.

fix input size index
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant