Skip to content

Instantly share code, notes, and snippets.

@shaypal5
Last active January 18, 2024 16:20
Show Gist options
  • Star 62 You must be signed in to star a gist
  • Fork 16 You must be signed in to fork a gist
  • 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
@Znigneering
Copy link

Thank you!

@BrunoGomesCoelho
Copy link

Hey, thanks for this! I modified it to work for normalized cm as well, you can just add this:

    if normalize:
        confusion_matrix = confusion_matrix.astype('float') / confusion_matrix.sum(axis=1)[:, np.newaxis]
        print("Normalized confusion matrix")
    else:
        print('Confusion matrix, without normalization')

and them add the format option:

    fmt = '.2f' if normalize else 'd'

and add fmt=fmt in your try-catch :)

@Akashpatel579
Copy link

Thanks for sharing

@huntermaxfield
Copy link

I can't seem to figure out why my heat map is printing twice? here is my code:

y_pred_knn_raw = knn_model.predict(X_test) print(metrics.confusion_matrix(y_true = y_test, y_pred = y_pred_knn)) #print basic confusion matrix print('Accuracy = ', metrics.accuracy_score(y_true = y_test, y_pred = y_pred_knn)) #printing accurary print_confusion_matrix(metrics.confusion_matrix(y_true = y_test, y_pred = y_pred_knn), labels) #print cm heatmap

Hoping someone could help me see what i'm doing wrong?

@guifeliper
Copy link

I can't seem to figure out why my heat map is printing twice? here is my code:

y_pred_knn_raw = knn_model.predict(X_test) print(metrics.confusion_matrix(y_true = y_test, y_pred = y_pred_knn)) #print basic confusion matrix print('Accuracy = ', metrics.accuracy_score(y_true = y_test, y_pred = y_pred_knn)) #printing accurary print_confusion_matrix(metrics.confusion_matrix(y_true = y_test, y_pred = y_pred_knn), labels) #print cm heatmap

Hoping someone could help me see what i'm doing wrong?

Printing Twice for me too!

@alanaor
Copy link

alanaor commented Mar 29, 2019

Add a ';' after you call the function. Should only print once.

@jeonghyunwoo
Copy link

Thank you, very helpful

@pranjalchaubey
Copy link

Prints twice for no apparent reason.

@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