Training: Fix confusion matrix conversion

softawaregmbh · Nov 17, 2022 · b8c2372 · b8c2372
1 parent 8449b7a
commit b8c2372
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/src/2_training.py b/src/2_training.py
@@ -96,7 +96,12 @@ def analyze_model(digit_classifier, x_test, y_test):
         confusion_matrix=confusion_matrix(y_test, y_pred, labels=digit_classifier.classes_),
         display_labels=digit_classifier.classes_
     )
-    mlflow.log_image(matplotlib_figure_to_pillow_image(confusion_matrix_display.figure_), 'test_confusion_matrix.png')
+    confusion_matrix_display.plot()
+
+    mlflow.log_image(
+        matplotlib_figure_to_pillow_image(confusion_matrix_display.figure_),
+        'test_confusion_matrix.png'
+    )
 
     end_action(action_text)
 

diff --git a/src/utils.py b/src/utils.py
@@ -1,5 +1,5 @@
 import json
-from matplotlib import figure
+from matplotlib import figure as Figure
 from PIL import Image as PILImage
 import subprocess
 import time
@@ -57,5 +57,7 @@ def request_user_consent(question: str) -> bool:
     return len(response) == 0 or response.lower() == 'y'
 
 
-def matplotlib_figure_to_pillow_image(figure: figure) -> PILImage:
+def matplotlib_figure_to_pillow_image(figure: Figure, not_drawn_before: bool = True) -> PILImage:
+    if not_drawn_before:
+        figure.canvas.draw()
     return PILImage.frombytes('RGB', figure.canvas.get_width_height(), figure.canvas.tostring_rgb())