tensorflow

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


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)?


Solution

  • 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();