Skip to content

Instantly share code, notes, and snippets.

@shaypal5
Last active May 13, 2024 13:04
Show Gist options
  • Save shaypal5/94c53d765083101efc0240d776a23823 to your computer and use it in GitHub Desktop.
Save shaypal5/94c53d765083101efc0240d776a23823 to your computer and use it in GitHub Desktop.
Pretty print a confusion matrix with seaborn
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
def print_confusion_matrix(confusion_matrix, class_names, figsize = (10,7), fontsize=14):
"""Prints a confusion matrix, as returned by sklearn.metrics.confusion_matrix, as a heatmap.
Note that due to returning the created figure object, when this funciton is called in a
notebook the figure willl be printed twice. To prevent this, either append ; to your
function call, or modify the function by commenting out the return expression.
Arguments
---------
confusion_matrix: numpy.ndarray
The numpy.ndarray object returned from a call to sklearn.metrics.confusion_matrix.
Similarly constructed ndarrays can also be used.
class_names: list
An ordered list of class names, in the order they index the given confusion matrix.
figsize: tuple
A 2-long tuple, the first value determining the horizontal size of the ouputted figure,
the second determining the vertical size. Defaults to (10,7).
fontsize: int
Font size for axes labels. Defaults to 14.
Returns
-------
matplotlib.figure.Figure
The resulting confusion matrix figure
"""
df_cm = pd.DataFrame(
confusion_matrix, index=class_names, columns=class_names,
)
fig = plt.figure(figsize=figsize)
try:
heatmap = sns.heatmap(df_cm, annot=True, fmt="d")
except ValueError:
raise ValueError("Confusion matrix values must be integers.")
heatmap.yaxis.set_ticklabels(heatmap.yaxis.get_ticklabels(), rotation=0, ha='right', fontsize=fontsize)
heatmap.xaxis.set_ticklabels(heatmap.xaxis.get_ticklabels(), rotation=45, ha='right', fontsize=fontsize)
plt.ylabel('True label')
plt.xlabel('Predicted label')
# Note that due to returning the created figure object, when this funciton is called in a notebook
# the figure willl be printed twice. To prevent this, either append ; to your function call, or
# modify the function by commenting out this return expression.
return fig
@laubosslink
Copy link

Prints twice for no apparent reason.

Just remove return fig

@nkc512
Copy link

nkc512 commented May 5, 2020

Thank you

@hector-margarito
Copy link

Awesome

@RunningJingJing
Copy link

awesome, thanks.
BTW "class_names" sort alphabetically, when there are a lot of attributes it will be better to do this as @scott Boston said
https://stackoverflow.com/questions/54875846/how-to-print-labels-and-column-names-for-confusion-matrix

and the argument normalize{‘true’, ‘pred’, ‘all’}, default=None can deal with the normalization when we generate the confusion matrix @BrunoGomesCoelho
eg:
confusion_matrix_array = confusion_matrix(ture_label, predict_label, normalize='all' )

@Arifur-R
Copy link

THANK YOU DEAR

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment