Stitching non max suppression (NMS) to YOLOv8n on exported ONNX model

Following my previous post on exploring YOLOv8, I have been stuck at using the YOLOv8 model other than PyTorch, because the direct export model give result of dimension like [batch size, 5, 8400], which does encapsulated the result of overlapped bounding boxes and confidence score. TF Lite with object detection API) would require post process of this result into bounding boxes that's not overlapped and corresponding confidence scores. As I observed, the YOLO class is initialized with member "model", which is the core model that would output that [batch size, 5, 8400] liked result. And the forward function is calling a predict function (that's depends on what task, be it object detection, classification or segmentation) Given the code's export towards Tensorflow related model (TF Lite, Tensorflow.js…) are through ONNX - Tensorflow saved model (through onnx2tf package) and then to the target. So ONNX seems to be a good target to add the necessary nms operation, a second reason would be the ONNX nms operation could be more optimized (compare to torchvision's nms converted operation).