Satisfying MyPy

TorchDSP · Sep 18, 2023 · 52297cb · 52297cb
1 parent b6e9b7c
commit 52297cb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git 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
@@ -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