Image classification problem with multiple right classes

Question

I have a use case where the model needs to detect fabricdefects. There are 15+ different kinds of defects. In one image there can be multiple defects present. The straight forward solution for this should be a multilabel model from my understanding. The classification of the data for a multilabel model is extremely tedious and errorprone.

Now I use a multiclass model, which seems to produce ok results. A multiclass model has only one right output class. My goal is to add one image to multiple classes and if the model predicts one of these classes it should result in a correct prediction (or lower loss output).

For example if an image contains defect1 and defect2 the model should look at both these outputs and calculate the loss from the one that has the highest probabilty.

Now my question: "Is it possible to have a model where the samples can have multiple right output classes, but are not penalized by predicting only one right class?"

Answer 1

This is a good question. There are definitely good reasons for wanting a loss function that evaluates whether at least one of the classes was picked up by the model.

To do what you are attempting, I believe you will have to write a custom loss function because this is sort of a "special request" kind of thing. I do not think any built-in loss functions in Tensorflow/Keras can do this, and I'm not really adept with other ML packages.

This how such a custom loss function could potentially work:

1. Once you have set a threshold for calling positive for each of the classes, you can convert your probability vector to a prediction vector of $1's$ and $0's$.

2. Then, you can do an elementwise check for equality to the ground truth (y_pred == y_true) and sum over all classes for which the defect is actually present (i.e., include $1 = 1$ but not $0 = 0$ in your tally).

3. If the sum is $>= 1$, then you know that the model found at least one defect correctly ($y_{pred} = 1$). Otherwise, the model did not predict the presence of any defects correctly ($y_{pred} = 0$).

4. You will have to sum over your original $y_{pred}$ vector, and if the sum $> =1$ (at least $1$ defect present), then the new $y_{pred} = 1$; otherwise, $y_{pred} = 0$.

5. Now, you have a new binary classification problem: Predict if defect is present or absent. You should be able to use binary crossentropy loss or another suitable loss function.

Here is a helpful page to get you started on custom loss functions in Keras. Secondly, I'm sure you know this already, but it will be important to make sure that your classes are balanced after this "data transformation". Finally, you can set up this problem a little differently such that you have 3 possible outputs: 1) correctly predicts at least 1 defect, 2) correctly predicts the absence of defects, and 3) makes no correct predicts.

Good luck. Hope that helps.