make colorbar optional in plot_confusion_matrix() reiinakano#114

Author: celik-muhammed

scikitplot/metrics.py

@@ -32,7 +32,7 @@ def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
                           pred_labels=None, title=None, normalize=False,
                           hide_zeros=False, hide_counts=False, x_tick_rotation=0, ax=None,
                           figsize=None, cmap='Blues', title_fontsize="large",
-                          text_fontsize="medium"):
+                          text_fontsize="medium", show_colorbar=True):
     """Generates confusion matrix plot from predictions and true labels
 
     Args:
@@ -89,6 +89,9 @@ def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
             Use e.g. "small", "medium", "large" or integer-values. Defaults to
             "medium".
 
+        show_colorbar (bool, optional): If False, does not add colour bar.
+            Defaults to True.
+
     Returns:
         ax (:class:`matplotlib.axes.Axes`): The axes on which the plot was
             drawn.
@@ -151,7 +154,10 @@ def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
         ax.set_title('Confusion Matrix', fontsize=title_fontsize)
 
     image = ax.imshow(cm, interpolation='nearest', cmap=plt.cm.get_cmap(cmap))
-    plt.colorbar(mappable=image)
+
+    if show_colorbar == True:
+        plt.colorbar(mappable=image)
+        
     x_tick_marks = np.arange(len(pred_classes))
     y_tick_marks = np.arange(len(true_classes))
     ax.set_xticks(x_tick_marks)