## Summary
This PR completes the TFLite `NON_MAX_SUPPRESSION_V5` implementation in
Relax by adding support for `soft_nms_sigma != 0`.
It extends `relax.vision.non_max_suppression` with soft-NMS attributes,
updates the TFLite frontend to consume the soft-NMS outputs correctly,
and aligns the TOPI implementation with LiteRT's reference behavior.
Relates to #19412.
## Changes
1. **Relax / TOPI soft-NMS support**
- Extend `NonMaximumSuppressionAttrs` with `soft_nms_sigma` and
`score_threshold`.
- Thread the new attributes through Relax op registration, Python
wrapper, and legalization.
- Add soft-NMS handling to TOPI classic NMS so
`relax.vision.non_max_suppression` can represent the
`NON_MAX_SUPPRESSION_V5` behavior.
2. **TFLite frontend support for `NON_MAX_SUPPRESSION_V5`**
- Remove the previous `soft_nms_sigma != 0` unsupported-path guard in
the TFLite frontend.
- Forward `soft_nms_sigma` and `score_threshold` into
`relax.vision.non_max_suppression`.
- Handle the soft-NMS return path explicitly so the frontend reads
decayed scores from the processed NMS output instead of re-reading the
original score tensor.
3. **Soft-NMS correctness fixes**
- Fix the soft-NMS path so boxes whose scores fall below the threshold
after decay are invalidated consistently.
- Keep returned indices and decayed scores aligned in both the TOPI TIR
implementation and the NumPy reference implementation.
- Update the soft-NMS candidate selection logic to re-pick the current
best candidate after each decay step, matching LiteRT's
reference behavior.
- Align the Gaussian decay formula with LiteRT.
4. **Test coverage**
- Add Relax tests for soft-NMS struct-info inference and legalization.
- Add Relax E2E tests covering reordered outputs after score decay and
other soft-NMS follow-up cases.
- Add TFLite frontend tests for `NON_MAX_SUPPRESSION_V5` with
`soft_nms_sigma != 0`.
- Add IR checks to verify that `soft_nms_sigma` and `score_threshold`
are forwarded correctly.
## Testing
```bash
python -m pytest -n 1 tests/python/relax/test_op_vision.py -k "all_class_non_max_suppression or get_valid_counts or nms" -v
python -m pytest tests/python/relax/test_frontend_tflite.py -k "nms_v5" -v
```
## Result:
- Relax vision tests passed locally
- TFLite `NON_MAX_SUPPRESSION_V5` coverage added for both hard-NMS and
soft-NMS paths
## Summary
Add `relax.vision.get_valid_counts` and classic
`relax.vision.non_max_suppression` for object-detection post-processing
pipelines.
`get_valid_counts` performs score-based bounding box filtering and
compacts valid boxes to the front of each batch. Classic
`non_max_suppression` performs flexible IoU-based suppression on
filtered boxes, complementing existing `all_class_non_max_suppression`
for custom post-processing workflows.
This PR implements the Relax-level registration, legalization, TOPI
compute, and test coverage for both operators.
## Changes
**Relax op registration and legalization:**
- C++ op functions, FFI registration, and struct info inference for both
operators (`vision.h`, `vision.cc`)
- Python wrappers with Relax docstrings (`vision.py`)
- Legalization to `topi.vision.get_valid_counts` and
`topi.vision.non_max_suppression`
- Additional struct-info validation for `score_index`, `id_index`, and
`coord_start` when `elem_length` is statically known
**TOPI and testing:**
- Full TOPI implementation for `get_valid_counts`
- Reimplementation of classic `non_max_suppression` in TOPI
- NumPy reference implementations in `tvm.topi.testing` for both
operators
- Op-level tests for struct info inference, legalization, invalid
attribute ranges, and e2e numerical correctness
- Stronger legalization tests that verify both `relax.call_tir`
introduction and removal of the original Relax vision op
## Limitations
- Attribute range validation for `score_index`, `id_index`, and
`coord_start` is only enforced when the input `elem_length` is
statically known during struct-info inference.
- Classic `non_max_suppression` follows the existing Relax / TOPI API
shape and is intended for single-class or class-aware custom
post-processing flows, distinct from `all_class_non_max_suppression`.
## Validation
```bash
pytest tests/python/relax/test_op_vision.py -k "get_valid_counts" -v
pytest tests/python/relax/test_op_vision.py -k "test_nms_" -v
```
All related tests passed.
ConvolutedDog