Created
February 10, 2020 21:02
-
-
Save eggie5/f5fa3e159e007a89494fc745f2f7e219 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def plot_confusion_matrix(y_true, y_pred, classes, | |
normalize=False, | |
title=None, | |
cmap=plt.cm.Blues): | |
""" | |
This function prints and plots the confusion matrix. | |
Normalization can be applied by setting `normalize=True`. | |
""" | |
if not title: | |
if normalize: | |
title = 'Normalized confusion matrix' | |
else: | |
title = 'Confusion matrix, without normalization' | |
# Compute confusion matrix | |
cm = confusion_matrix(y_true, y_pred) | |
# Only use the labels that appear in the data | |
classes = classes[unique_labels(y_true, y_pred)] | |
if normalize: | |
cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] | |
print("Normalized confusion matrix") | |
else: | |
print('Confusion matrix, without normalization') | |
print(cm) | |
fig, ax = plt.subplots() | |
im = ax.imshow(cm, interpolation='nearest', cmap=cmap) | |
ax.figure.colorbar(im, ax=ax) | |
# We want to show all ticks... | |
ax.set(xticks=np.arange(cm.shape[1]), | |
yticks=np.arange(cm.shape[0]), | |
# ... and label them with the respective list entries | |
xticklabels=classes, yticklabels=classes, | |
title=title, | |
ylabel='True label', | |
xlabel='Predicted label') | |
# Rotate the tick labels and set their alignment. | |
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", | |
rotation_mode="anchor") | |
# Loop over data dimensions and create text annotations. | |
fmt = '.2f' if normalize else 'd' | |
thresh = cm.max() / 2. | |
for i in range(cm.shape[0]): | |
for j in range(cm.shape[1]): | |
ax.text(j, i, format(cm[i, j], fmt), | |
ha="center", va="center", | |
color="white" if cm[i, j] > thresh else "black") | |
fig.tight_layout() | |
return ax |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment