How to create a manual confusion matrix using TensorFlow.js?

Question

I am trying to create a confusion matrix using Tensorflow.js which shows True Positive, False Positive, True Negative and False Negative.

Here is the script to reproduce: https://github.com/erostefano/humanSunglassesAnalyzer.js/blob/feature/training/training/src/index.js.

I could do the mapping by myself but I am not getting the correct values. Actually I want to compare the return of cnn.predict(xTest) with my labels (yTest).

How do I get the predicted labels from cnn.predict(xTest)?

Answer 1

The labels can be calculated using argmax:

const predictions = cnn.predict(xTest);
const predictedLabels = predictions.argMax(-1).dataSync();

Then you can calculate the true labels:

const trueLabels = yTest.argMax(-1).dataSync();

The confusion matrix can then easily calculated like this:

const numClasses = yTest.shape[1];
const confusionMatrix = tf.math.confusionMatrix(trueLabels, predictedLabels, numClasses);
confusionMatrix.print();