pythonmatplotlibscikit-learnconfusion-matrix

Show empty rows and columns in sklearn ConfusionMatrixDisplay.from_predictions


I want to display a confusion matrix on label prediction. On certain subsets of my data, some classes are missing (from both the ground truth and prediction), eg class 6 in the example below.

I would like to still show the 6 line and column even if filled with zeros. Is there a way to achieve this with sklearn's ConfusionMatrixDisplay ?

from sklearn.metrics import ConfusionMatrixDisplay
import matplotlib.pyplot as plt


GT   = [4, 4, 4, 5, 2, 1, 1, 1, 2, 7]
pred = [4, 5, 4, 5, 3, 0, 3, 2, 2, 7]

fig, ax = plt.subplots()
ConfusionMatrixDisplay.from_predictions(GT, pred,
                                        cmap='Blues',
                                        ax=ax,)
plt.show()

I have tried to pass my list of labels to the display_labels argument, but it fails because the length does not match the number of classes that is actually in the data


GT   = [4, 4, 4, 5, 2, 1, 1, 1, 2, 7]
pred = [4, 5, 4, 5, 3, 0, 3, 2, 2, 7]

fig, ax = plt.subplots()
ConfusionMatrixDisplay.from_predictions(GT, pred,
                                        cmap='Blues',
                                        ax=ax,
                                        display_labels = range(8))
plt.show()

enter image description here


Solution

  • The from_estimator and from_predictions methods have a labels argument to set the labels directly, overriding what's in the data. See the full list of options in the docs page.