mypy fixes

torchgeo/datasets/utils.py

@@ -433,7 +433,7 @@ def draw_semantic_segmentation_masks(
     mask: Tensor,
     alpha: float = 0.5,
     colors: Optional[ColorMap] = None,
-) -> np.ndarray:
+) -> np.ndarray:  # type: ignore[type-arg]
     """Overlay a semantic segmentation mask onto an image.
 
     Args:
@@ -444,11 +444,11 @@ def draw_semantic_segmentation_masks(
     Returns:
         a list of the subset datasets. Either [train, val] or [train, val, test]
     """
-    classes = torch.unique(mask)
+    classes = torch.unique(mask)  # type: ignore[attr-defined]
     classes = classes[1:]
     class_masks = mask == classes[:, None, None]
     img = draw_segmentation_masks(
         image=image, masks=class_masks, alpha=alpha, colors=colors
     )
-    img = img.permute((1, 2, 0))
-    return img.numpy()
+    img = img.permute((1, 2, 0)).numpy()
+    return img  # type: ignore[no-any-return]

torchgeo/datasets/xview.py

@@ -182,10 +182,16 @@ def plot(self, index: int, alpha: float = 0.5) -> plt.Figure:
         """
         sample = self[index]
         image1 = draw_semantic_segmentation_masks(
-            sample["image"][0], sample["mask"][0], alpha=alpha, colors=self.colormap
+            sample["image"][0],
+            sample["mask"][0],
+            alpha=alpha,
+            colors=self.colormap  # type: ignore[arg-type]
         )
         image2 = draw_semantic_segmentation_masks(
-            sample["image"][1], sample["mask"][1], alpha=alpha, colors=self.colormap
+            sample["image"][1],
+            sample["mask"][1],
+            alpha=alpha,
+            colors=self.colormap  # type: ignore[arg-type]
         )
         fig, (ax1, ax2) = plt.subplots(ncols=2)
         fig.set_size_inches((25, 25))