Fix some issues raised by mypy

Author: constantinpape

micro_sam/instance_segmentation.py

@@ -1,10 +1,9 @@
 import multiprocessing as mp
 import warnings
 from abc import ABC
-from collections.abc import Mapping
 from concurrent import futures
 from copy import deepcopy
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 import torch
@@ -46,7 +45,7 @@ def __getitem__(self, index):
 
 
 def mask_data_to_segmentation(
-    masks: List[Mapping[str, Any]],
+    masks: List[Dict[str, Any]],
     shape: tuple[int, ...],
     with_background: bool,
 ) -> np.ndarray:
@@ -256,7 +255,7 @@ def _postprocess_masks(self, mask_data, min_mask_region_area, box_nms_thresh, cr
 
         return curr_anns
 
-    def get_state(self) -> Mapping[str, Any]:
+    def get_state(self) -> Dict[str, Any]:
         """Get the initialized state of the mask generator.
 
         Returns:
@@ -266,7 +265,7 @@ def get_state(self) -> Mapping[str, Any]:
             raise RuntimeError("The state has not been computed yet. Call initialize first.")
         return {"crop_list": self.crop_list, "crop_boxes": self.crop_boxes, "original_size": self.original_size}
 
-    def set_state(self, state: Mapping[str, Any]) -> None:
+    def set_state(self, state: Dict[str, Any]) -> None:
         """Set the state of the mask generator.
 
         Args:
@@ -447,7 +446,7 @@ def generate(
         crop_nms_thresh: float = 0.7,
         min_mask_region_area: int = 0,
         output_mode: str = "binary_mask",
-    ) -> List[Mapping[str, Any]]:
+    ) -> List[Dict[str, Any]]:
         """Generate instance segmentation for the currently initialized image.
 
         Args:
@@ -486,7 +485,7 @@ def generate(
                 data["boxes"].float(),
                 scores,
                 torch.zeros_like(data["boxes"][:, 0]),  # categories
-                iou_threshold=self.crop_nms_thresh,
+                iou_threshold=crop_nms_thresh,
             )
             data.filter(keep_by_nms)
 
@@ -648,7 +647,7 @@ def generate(
         box_nms_thresh: float = 0.7,
         min_mask_region_area: int = 0,
         output_mode: str = "binary_mask",
-    ) -> List[Mapping[str, Any]]:
+    ) -> List[Dict[str, Any]]:
         """Generate instance segmentation for the currently initialized image.
 
         Args:
@@ -699,7 +698,7 @@ def get_initial_segmentation(self) -> np.ndarray:
             raise RuntimeError("AutomaticMaskGenerator has not been initialized. Call initialize first.")
         return self._resize_segmentation(self._initial_segmentation, self.original_size)
 
-    def get_state(self) -> Mapping[str, Any]:
+    def get_state(self) -> Dict[str, Any]:
         """Get the initialized state of the mask generator.
 
         Returns:
@@ -709,7 +708,7 @@ def get_state(self) -> Mapping[str, Any]:
         state["initial_segmentation"] = self._initial_segmentation
         return state
 
-    def set_state(self, state: Mapping[str, Any]) -> None:
+    def set_state(self, state: Dict[str, Any]) -> None:
         """Set the state of the mask generator.
 
         Args:
@@ -918,7 +917,7 @@ def segment_tile(_, tile_id):
 
         return segmentation
 
-    def get_initial_segmentation(self) -> None:
+    def get_initial_segmentation(self) -> np.ndarray:
         """Get the initial instance segmentation.
 
         Returns:
@@ -947,7 +946,7 @@ def segment_tile(_, tile_id):
         self._stitched_initial_segmentation = initial_segmentation
         return initial_segmentation
 
-    def get_state(self) -> Mapping[str, Any]:
+    def get_state(self) -> Dict[str, Any]:
         """Get the initialized state of the mask generator.
 
         Returns:
@@ -958,7 +957,7 @@ def get_state(self) -> Mapping[str, Any]:
         state["halo"] = self._halo
         return state
 
-    def set_state(self, state: Mapping[str, Any]) -> None:
+    def set_state(self, state: Dict[str, Any]) -> None:
         """Set the state of the mask generator.
 
         Args:

micro_sam/prompt_generators.py

@@ -37,7 +37,7 @@ def __init__(
         dilation_strength: int,
         get_point_prompts: bool = True,
         get_box_prompts: bool = False
-    ):
+    ) -> None:
         self.n_positive_points = n_positive_points
         self.n_negative_points = n_negative_points
         self.dilation_strength = dilation_strength

micro_sam/util.py

@@ -1,10 +1,10 @@
 import hashlib
 import os
 import warnings
-from collections.abc import Mapping
 from shutil import copyfileobj
-from typing import Any, Optional
+from typing import Any, Callable, Dict, Optional, Tuple
 
+import imageio.v3 as imageio
 import numpy as np
 import requests
 import torch
@@ -17,11 +17,6 @@
 
 from segment_anything import sam_model_registry, SamPredictor
 
-try:
-    import imageio.v2 as imageio
-except ImportError:
-    import imageio
-
 try:
     from napari.utils import progress as tqdm
 except ImportError:
@@ -339,9 +334,9 @@ def precompute_image_embeddings(
     save_path: Optional[str] = None,
     lazy_loading: bool = False,
     ndim: Optional[int] = None,
-    tile_shape: Optional[tuple[int]] = None,
-    halo: Optional[tuple[int]] = None,
-    wrong_file_callback: Optional[callable] = None,
+    tile_shape: Optional[Tuple[int]] = None,
+    halo: Optional[Tuple[int]] = None,
+    wrong_file_callback: Optional[Callable] = None,
 ) -> ImageEmbeddings:
     """Compute the image embeddings (output of the encoder) for the input.
 
@@ -453,7 +448,7 @@ def compute_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
 def get_centers_and_bounding_boxes(
     segmentation: np.ndarray,
     mode: str = "v"
-) -> tuple[Mapping[int, np.ndarray], Mapping[int, tuple]]:
+) -> Tuple[Dict[int, np.ndarray], Dict[int, tuple]]:
     """Returns the center coordinates of the foreground instances in the ground-truth.
 
     Args:
@@ -500,7 +495,7 @@ def load_image_data(
         The image data.
     """
     if key is None:
-        image_data = imageio.imread(path) if ndim == 2 else imageio.volread(path)
+        image_data = imageio.imread(path)
     else:
         with open_file(path, mode="r") as f:
             image_data = f[key]