Plot confusion matrix in black and white

Question

I currently have the following code:

disp = ConfusionMatrixDisplay(confusion_matrix=cm)
disp.plot()
plt.show()

This results into something like:

However, I want the diagonal to be depicted with a black background, and the off-diagonal with a white background, that is something like:

How can I change the cell colors accordingly?

Answer 1

You can create a custom colormap where index 0 = white and index 1 = black like so:

cmap = ListedColormap(['white', 'black'])
color_matrix = np.eye(cm.shape[0], dtype=int)

When imshow() uses this color_matrix with the cmap it maps:

• 0 → white (off-diagonal/incorrect predictions)

• 1 → black (diagonal/correct predictions)

Reference

The code:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import ListedColormap
from sklearn.metrics import ConfusionMatrixDisplay
cm = np.array([[12, 1],
             [2, 10]])
def plot_bw_cm(cm, labels=None):
  cmap = ListedColormap(['white', 'black'])
  color_matrix = np.eye(cm.shape[0], dtype=int)
  fig, ax = plt.subplots(figsize=(6, 5))
  im = ax.imshow(color_matrix, cmap=cmap, aspect='equal', vmin=0, vmax=1)
  
  for i in range(cm.shape[0]):
      for j in range(cm.shape[1]):
          text_color = 'white' if i == j else 'black'
          ax.text(j, i, str(cm[i, j]), ha='center', va='center',
                 color=text_color, fontsize=16, fontweight='bold')
  
  if labels is None:
      labels = range(cm.shape[0])
  ax.set_xticks(range(cm.shape[1]))
  ax.set_yticks(range(cm.shape[0]))
  ax.set_xticklabels(labels)
  ax.set_yticklabels(labels)
  ax.set_xlabel('Predicted label')
  ax.set_ylabel('True label')
  
  plt.tight_layout()
  plt.show()

plot_bw_cm(cm)