How to limit YOLO to a fixed number of predictions?

Question

Let's say I annotated all images in my dataset to have 20 Bounding Boxes.

I basically want my predicted bounding boxes to also be only 20. After training however, I get differing amounts of bounding boxes, that aren't 20.

I'm trying to detect the same 20 objects in an image. All the objects are the same so I only 1 class for all 20 bounding boxes.

I'm currently using YOLOv5 but is there a better model for a use-case like this?

Answer 1

I suggest selecting the 20 detected objects with higher confidence, you can do that easily by appending the all detected objects boxes to a list as well as the confidence and labels and then iterated through the list with range limitation, and then you can draw the bounding box of the filtered objects (20 objects).