From 52297cb9ab0a3fcce71f9d65c71bb4a82edb340c Mon Sep 17 00:00:00 2001 From: Garrett Vanhoy Date: Mon, 18 Sep 2023 10:59:59 -0400 Subject: [PATCH] Satisfying MyPy --- torchsig/utils/cm_plotter.py | 3 ++- torchsig/utils/visualize.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/torchsig/utils/cm_plotter.py b/torchsig/utils/cm_plotter.py index 75969b7..34c74f6 100644 --- a/torchsig/utils/cm_plotter.py +++ b/torchsig/utils/cm_plotter.py @@ -1,5 +1,6 @@ from sklearn.metrics import confusion_matrix from matplotlib import pyplot as plt +from matplotlib.colors import Colormap from typing import Optional import numpy as np @@ -13,7 +14,7 @@ def plot_confusion_matrix( text: bool = True, rotate_x_text: int = 90, figsize: tuple = (16, 9), - cmap: plt.cm = plt.cm.Blues, + cmap: str = "Blues", ): """Function to help plot confusion matrices diff --git a/torchsig/utils/visualize.py b/torchsig/utils/visualize.py index b4512aa..dede717 100644 --- a/torchsig/utils/visualize.py +++ b/torchsig/utils/visualize.py @@ -176,7 +176,7 @@ def _visualize(self, iq_data: np.ndarray, targets: np.ndarray) -> Figure: ts = np.arange(len(cwt_matrix[0])) / self.sample_rate plt.imshow( np.abs(cwt_matrix), - extent=[ts[0], ts[-1], freqs[-1], freqs[0]], + extent=[ts[0], ts[-1], freqs[-1], freqs[0]], # type: ignore vmin=0, vmax=np.abs(cwt_matrix).max(), aspect="auto", @@ -458,7 +458,7 @@ def _visualize( # type: ignore title = "Data" plt.xticks([]) plt.yticks([]) - plt.title(title) + plt.title(str(title)) return figure @@ -521,7 +521,7 @@ def _visualize(self, data: np.ndarray, targets: np.ndarray) -> Figure: title = "Data" plt.xticks([]) plt.yticks([]) - plt.title(title) + plt.title(str(title)) return figure