Detection model - Training with class-instance limit-awareness

Question

Is there a way to make a detection model aware of the maximal number of possible objects of a given class within a single image?

For example in a toy case with 2 classes. If I know that in every single image, class A can have no more than 5 instances and class B no more than 1. Is there a way to incorporate it into the training process?

To make it clear, I'm not talking about an additional algorithm which runs on top of the trained model (such as non-maximum suppression which is used to select a single bounding box for an object). I specifically ask about the actual model and its training process.

Answer 1

I can think of the following approach:

Let's say that you have two classes, A and B. Additionally, you now that for class A there is at maximum 5 instances (so 0, 1, 2, 3, 4 or 5) and for B 1 instance (0 or 1).

For this purpose, you can have 6 outputs for class A and 2 outputs for class B. Between those 6 outputs for class A, only one should be active; same for B - only one of two should be active.

For example, if on some image there are 3 objects of class A and 0 for B, the outputs would be: [0, 0, 0, 1, 0, 0] for A, and [1, 0] for B (or something very close to 0s and 1s, right?)

With these outputs, you can also combine other outputs which are needed for detection.

Answer 2

This is a very interesting supervision but hard to achieve!

Why we need this supervision?

The need for this supervision comes from the fact that model may wrongly detect more objects than it should, thus, must be punished (taught) for this violation, otherwise no supervision would be required since model is acting accordingly.

How to implement this supervision?

To this end, we need to fork some layers from the model to output the number of detected objects per class $c$ for input image $i$, namely $n'_{c,i}$, then supervise this output with the true number of objects in image $i$, namely $n_{c, i}$, or merely with an upper limit $N_c$ per class as you have suggested. Then, add a term like $(n_{c, i} - n'_{c, i})^2$ or $\text{max}(0, n'_{c, i} - N_c)$ to the loss function to punish the model for detecting wrong or more number of objects than it should respectively. Then proceed to train the model.

What may go wrong?

But here is the problem, model can learn to lie about the number of detected objects through modifying the forked layers (weights)! Since it is easier for model to fabricate a valid $n'_{c, i}$ than to actually detect fewer objects which is more complex. Also, if we use a constant, unfabricatable unit (e.g. a constant neural net) that counts the number of detected objects, there would be no gradient to punish (teach) the model!

This is why this type of supervision is hard to achieve.