search

Deci-AI / super-gradients

Easily train or fine-tune SOTA computer vision models with one open source training library. The home of Yolo-NAS.
https://www.supergradients.com
Apache License 2.0
4.43k stars 482 forks source link

Enable class agnostic NMS #1985

Open kostastsing opened 2 months ago

kostastsing commented 2 months ago

💡 Your Question

Any idea on how to enable class agnostic NMS when exporting to onnx using TensorRT backend? There is a class_agnostic parameter for the EfficientNMS_TRT (see here) plugin but I haven't figured out yet how to set it to True

Versions

3.7.0

BloodAxe commented 2 months ago

Hi there. Great question!

Out of the box this is not supported, but it is actually not that hard to implement.

First step is to set this attribute in attach_tensorrt_nms method and make this an attribute of a function.

Second step would be to propagate this variable up to the model.export so that is is accessible from the public export API.

We would be happy to get a PR with this enhancement 😎 if you manage to get it working :)