🧑‍💻 Fix mypy Type Error in tiatoolbox/tools (#842)

- Revisited the files in tools/, and fixed the errors that were not detected previously.

Author: Jiaqi-Lv

tiatoolbox/tools/graph.py

@@ -4,7 +4,7 @@
 
 from collections import defaultdict
 from numbers import Number
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Callable, cast
 
 import numpy as np
 import torch
@@ -18,15 +18,15 @@
     from numpy.typing import ArrayLike
 
 
-def delaunay_adjacency(points: ArrayLike, dthresh: float) -> list:
+def delaunay_adjacency(points: np.ndarray, dthresh: float) -> np.ndarray:
     """Create an adjacency matrix via Delaunay triangulation from a list of coordinates.
 
     Points which are further apart than dthresh will not be connected.
 
     See https://en.wikipedia.org/wiki/Adjacency_matrix.
 
     Args:
-        points (ArrayLike):
+        points (np.ndarray):
             An nxm list of coordinates.
         dthresh (float):
             Distance threshold for triangulation.
@@ -113,11 +113,11 @@ def triangle_signed_area(triangle: ArrayLike) -> int:
     )
 
 
-def edge_index_to_triangles(edge_index: ArrayLike) -> ArrayLike:
+def edge_index_to_triangles(edge_index: np.ndarray) -> np.ndarray:
     """Convert an edged index to triangle simplices (triplets of coordinate indices).
 
     Args:
-        edge_index (ArrayLike):
+        edge_index (np.ndarray):
             An Nx2 array of edges.
 
     Returns:
@@ -157,24 +157,24 @@ def edge_index_to_triangles(edge_index: ArrayLike) -> ArrayLike:
 
 
 def affinity_to_edge_index(
-    affinity_matrix: torch.Tensor | ArrayLike,
+    affinity_matrix: torch.Tensor | np.ndarray,
     threshold: float = 0.5,
-) -> torch.tensor | ArrayLike:
+) -> torch.Tensor | np.ndarray:
     """Convert an affinity matrix (similarity matrix) to an edge index.
 
     Converts an NxN affinity matrix to a 2xM edge index, where M is the
     number of node pairs with a similarity greater than the threshold
     value (defaults to 0.5).
 
     Args:
-        affinity_matrix:
+        affinity_matrix (torch.Tensor | np.ndarray):
             An NxN matrix of affinities between nodes.
         threshold (Number):
             Threshold above which to be considered connected. Defaults
             to 0.5.
 
     Returns:
-        ArrayLike or torch.Tensor:
+        torch.Tensor | np.ndarray:
             The edge index of shape (2, M).
 
     Example:
@@ -191,9 +191,9 @@ def affinity_to_edge_index(
         raise ValueError(msg)
     # Handle cases for pytorch and numpy inputs
     if isinstance(affinity_matrix, torch.Tensor):
-        return (affinity_matrix > threshold).nonzero().t().contiguous()
+        return (affinity_matrix > threshold).nonzero().t().contiguous().to(torch.int64)
     return np.ascontiguousarray(
-        np.stack((affinity_matrix > threshold).nonzero(), axis=1).T,
+        np.stack((affinity_matrix > threshold).nonzero(), axis=1).T.astype(np.int64),
     )
 
 
@@ -208,7 +208,7 @@ class SlideGraphConstructor:
     """
 
     @staticmethod
-    def _umap_reducer(graph: dict[str, ArrayLike]) -> ArrayLike:
+    def _umap_reducer(graph: dict[str, np.ndarray]) -> np.ndarray:
         """Default reduction which reduces `graph["x"]` to 3D values.
 
         Reduces graph features to 3D values using UMAP which are suitable
@@ -220,7 +220,7 @@ def _umap_reducer(graph: dict[str, ArrayLike]) -> ArrayLike:
                 "coordinates".
 
         Returns:
-            ArrayLike:
+            np.ndarray:
                 A UMAP embedding of `graph["x"]` with shape (N, 3) and
                 values ranging from 0 to 1.
         """
@@ -232,15 +232,15 @@ def _umap_reducer(graph: dict[str, ArrayLike]) -> ArrayLike:
 
     @staticmethod
     def build(
-        points: ArrayLike,
-        features: ArrayLike,
+        points: np.ndarray,
+        features: np.ndarray,
         lambda_d: float = 3.0e-3,
         lambda_f: float = 1.0e-3,
         lambda_h: float = 0.8,
         connectivity_distance: int = 4000,
         neighbour_search_radius: int = 2000,
         feature_range_thresh: float | None = 1e-4,
-    ) -> dict[str, ArrayLike]:
+    ) -> dict[str, np.ndarray]:
         """Build a graph via hybrid clustering in spatial and feature space.
 
         The graph is constructed via hybrid hierarchical clustering
@@ -266,10 +266,10 @@ def build(
         connected.
 
         Args:
-            points (ArrayLike):
+            points (np.ndarray):
                 A list of (x, y) spatial coordinates, e.g. pixel
                 locations within a WSI.
-            features (ArrayLike):
+            features (np.ndarray):
                 A list of features associated with each coordinate in
                 `points`. Must be the same length as `points`.
             lambda_d (Number):
@@ -400,27 +400,27 @@ def build(
             # Find the xy and feature space averages of the cluster
             point_centroids.append(np.round(points[idx, :].mean(axis=0)))
             feature_centroids.append(features[idx, :].mean(axis=0))
-        point_centroids = np.array(point_centroids)
-        feature_centroids = np.array(feature_centroids)
+        point_centroids_arr = np.array(point_centroids)
+        feature_centroids_arr = np.array(feature_centroids)
 
         adjacency_matrix = delaunay_adjacency(
-            points=point_centroids,
+            points=point_centroids_arr,
             dthresh=connectivity_distance,
         )
         edge_index = affinity_to_edge_index(adjacency_matrix)
-
+        edge_index = cast(np.ndarray, edge_index)
         return {
-            "x": feature_centroids,
-            "edge_index": edge_index.astype(np.int64),
-            "coordinates": point_centroids,
+            "x": feature_centroids_arr,
+            "edge_index": edge_index,
+            "coordinates": point_centroids_arr,
         }
 
     @classmethod
     def visualise(
         cls: type[SlideGraphConstructor],
-        graph: dict[str, ArrayLike],
-        color: ArrayLike | str | Callable | None = None,
-        node_size: Number | ArrayLike | Callable = 25,
+        graph: dict[str, np.ndarray],
+        color: np.ndarray | str | Callable | None = None,
+        node_size: int | np.ndarray | Callable = 25,
         edge_color: str | ArrayLike = (0, 0, 0, 0.33),
         ax: Axes | None = None,
     ) -> Axes:
@@ -510,7 +510,8 @@ def visualise(
 
         # Plot the nodes
         plt.scatter(
-            *nodes.T,
+            x=nodes.T[0],
+            y=nodes.T[1],
             c=color(graph) if callable(color) else color,
             s=node_size(graph) if callable(node_size) else node_size,
             zorder=2,

tiatoolbox/tools/patchextraction.py

@@ -45,7 +45,7 @@ class ExtractorParams(TypedDict, total=False):
     pad_mode: str
     pad_constant_values: int | tuple[int, int]
     within_bound: bool
-    input_mask: str | Path | np.ndarray | wsireader.WSIReader
+    input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader
     stride: int | tuple[int, int]
     min_mask_ratio: float
 
@@ -81,7 +81,7 @@ class SlidingWindowPatchExtractorParams(TypedDict):
     pad_mode: str
     pad_constant_values: int | tuple[int, int]
     within_bound: bool
-    input_mask: str | Path | np.ndarray | wsireader.WSIReader | None
+    input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader | None
     stride: int | tuple[int, int] | None
     min_mask_ratio: float
 
@@ -113,7 +113,8 @@ class PatchExtractor(PatchExtractorABC):
             Input image for patch extraction.
         patch_size(int or tuple(int)):
             Patch size tuple (width, height).
-        input_mask(str, pathlib.Path, :class:`numpy.ndarray`, or :obj:`WSIReader`):
+        input_mask
+            (str, pathlib.Path, :class:`numpy.ndarray`, or :obj:`VirtualWSIReader`):
             Input mask that is used for position filtering when
             extracting patches i.e., patches will only be extracted
             based on the highlighted regions in the input_mask.
@@ -187,7 +188,7 @@ def __init__(
         self: PatchExtractor,
         input_img: str | Path | np.ndarray,
         patch_size: int | tuple[int, int],
-        input_mask: str | Path | np.ndarray | wsireader.WSIReader | None = None,
+        input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader | None = None,
         resolution: Resolution = 0,
         units: Units = "level",
         pad_mode: str = "constant",
@@ -391,7 +392,7 @@ def filter_coordinates(
             0,
             tissue_mask.shape[0],
         )
-        scaled_coords = list((scaled_coords).astype(np.int32))
+        scaled_coords_list = list((scaled_coords).astype(np.int32))
 
         def default_sel_func(
             tissue_mask: np.ndarray,
@@ -412,7 +413,7 @@ def default_sel_func(
             ) and (pos_area > 0 and patch_area > 0)
 
         func = default_sel_func if func is None else func
-        flag_list = [func(tissue_mask, coord) for coord in scaled_coords]
+        flag_list = [func(tissue_mask, coord) for coord in scaled_coords_list]
 
         return np.array(flag_list)
 
@@ -529,7 +530,7 @@ def get_coordinates(
             msg = f"`stride_shape` value {stride_shape_arr} must > 1."
             raise ValueError(msg)
 
-        def flat_mesh_grid_coord(x: int, y: int) -> np.ndarray:
+        def flat_mesh_grid_coord(x: np.ndarray, y: np.ndarray) -> np.ndarray:
             """Helper function to obtain coordinate grid."""
             xv, yv = np.meshgrid(x, y)
             return np.stack([xv.flatten(), yv.flatten()], axis=-1)
@@ -577,7 +578,8 @@ class SlidingWindowPatchExtractor(PatchExtractor):
             Input image for patch extraction.
         patch_size(int or tuple(int)):
             Patch size tuple (width, height).
-        input_mask(str, pathlib.Path, :class:`numpy.ndarray`, or :obj:`WSIReader`):
+        input_mask
+            (str, pathlib.Path, :class:`numpy.ndarray`, or :obj:`VirtualWSIReader`):
             Input mask that is used for position filtering when
             extracting patches i.e., patches will only be extracted
             based on the highlighted regions in the `input_mask`.
@@ -627,7 +629,7 @@ def __init__(
         self: SlidingWindowPatchExtractor,
         input_img: str | Path | np.ndarray,
         patch_size: int | tuple[int, int],
-        input_mask: str | Path | np.ndarray | wsireader.WSIReader | None = None,
+        input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader | None = None,
         resolution: Resolution = 0,
         units: Units = "level",
         stride: int | tuple[int, int] | None = None,

tiatoolbox/tools/pyramid.py

@@ -223,7 +223,7 @@ def get_tile(
             output_size = self.output_tile_size // 2 ** (
                 self.sub_tile_level_count - level
             )
-            output_size = np.repeat(output_size, 2).astype(int)
+            output_size = np.repeat(output_size, 2).astype(int).tolist()
             thumb = self.get_thumb_tile()
             thumb.thumbnail((output_size[0], output_size[1]))
             return thumb
@@ -236,7 +236,7 @@ def get_tile(
         logger.addFilter(duplicate_filter)
         tile = self.wsi.read_rect(
             coord,
-            size=[v * res for v in output_size],
+            size=(output_size[0] * res, output_size[1] * res),
             resolution=res / scale,
             units="baseline",
             pad_mode=pad_mode,
@@ -510,7 +510,7 @@ class AnnotationTileGenerator(ZoomifyGenerator):
 
     """
 
-    def __init__(
+    def __init__(  # skipcq: PYL-W0231
         self: AnnotationTileGenerator,
         info: WSIMeta,
         store: AnnotationStore,
@@ -520,9 +520,11 @@ def __init__(
         overlap: int = 0,
     ) -> None:
         """Initialize :class:`AnnotationTileGenerator`."""
-        super().__init__(None, tile_size, downsample, overlap)
         self.info = info
         self.store = store
+        self.tile_size = tile_size
+        self.downsample = downsample
+        self.overlap = overlap
         if renderer is None:
             renderer = AnnotationRenderer()
         self.renderer = renderer

tiatoolbox/tools/registration/wsi_registration.py

@@ -1688,6 +1688,7 @@ def read_rect(
         transformed_patch = transformed_patch[start_row:end_row, start_col:end_col, :]
 
         # Resize to desired size
+        post_read_scale = float(post_read_scale[0]), float(post_read_scale[1])
         return imresize(
             img=transformed_patch,
             scale_factor=post_read_scale,

tiatoolbox/tools/stainaugment.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import copy
+from typing import cast
 
 import numpy as np
 from albumentations.core.transforms_interface import ImageOnlyTransform
@@ -132,7 +133,7 @@ def __init__(
 
         self.alpha: float
         self.beta: float
-        self.img_shape = None
+        self.img_shape: tuple[int, ...]
         self.tissue_mask: np.ndarray
         self.n_stains: int
         self.source_concentrations: np.ndarray
@@ -196,6 +197,7 @@ def augment(self: StainAugmentor) -> np.ndarray:
             else:
                 augmented_concentrations[self.tissue_mask, i] *= self.alpha
                 augmented_concentrations[self.tissue_mask, i] += self.beta
+        self.stain_matrix = cast(np.ndarray, self.stain_matrix)
         img_augmented = 255 * np.exp(
             -1 * np.dot(augmented_concentrations, self.stain_matrix),
         )

tiatoolbox/tools/tissuemask.py

@@ -124,7 +124,7 @@ def fit(
         # Convert RGB images to greyscale
         grey_images = [x[..., 0] for x in images]
         if images_shape[-1] == 3:  # noqa: PLR2004
-            grey_images = np.zeros(images_shape[:-1], dtype=np.uint8)
+            grey_images = np.zeros(images_shape[:-1], dtype=np.uint8).tolist()
             for n, image in enumerate(images):
                 grey_images[n] = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
 
@@ -206,7 +206,7 @@ def __init__(
         *,
         mpp: float | tuple[float, float] | None = None,
         power: float | tuple[float, float] | None = None,
-        kernel_size: int | tuple[int, int] | None = None,
+        kernel_size: int | tuple[int, int] | np.ndarray | None = None,
         min_region_size: int | None = None,
     ) -> None:
         """Initialise a morphological masker.
@@ -249,7 +249,7 @@ def __init__(
             mpp_array = np.array(mpp)
             if mpp_array.size != 2:  # noqa: PLR2004
                 mpp_array = mpp_array.repeat(2)
-            kernel_size = np.max([32 / mpp_array, [1, 1]], axis=0)
+            kernel_size = np.max([32 / mpp_array, np.array([1, 1])], axis=0)
 
         # Ensure kernel_size is a length 2 numpy array
         kernel_size_array = np.array(kernel_size)

tiatoolbox/utils/misc.py

@@ -482,12 +482,12 @@ def __assign_unknown_class(input_table: pd.DataFrame) -> pd.DataFrame:
 
 
 def read_locations(
-    input_table: PathLike | np.ndarray | pd.DataFrame,
+    input_table: str | Path | PathLike | np.ndarray | pd.DataFrame,
 ) -> pd.DataFrame:
     """Read annotations as pandas DataFrame.
 
     Args:
-        input_table (PathLike | np.ndarray | pd.DataFrame`):
+        input_table (str| Path| PathLike | np.ndarray | pd.DataFrame`):
             Path to csv, npy or json. Input can also be a
             :class:`numpy.ndarray` or :class:`pandas.DataFrame`.
             First column in the table represents x position, second

tiatoolbox/wsicore/wsireader.py

@@ -1136,7 +1136,7 @@ def read_rect(
         units: Units = "level",
         interpolation: str = "optimise",
         pad_mode: str = "constant",
-        pad_constant_values: Number | Iterable[NumPair] = 0,
+        pad_constant_values: int | tuple[int, int] = 0,
         coord_space: str = "baseline",
         **kwargs: dict,
     ) -> np.ndarray:
@@ -1185,7 +1185,7 @@ def read_rect(
                 Defaults to 'constant'. See :func:`numpy.pad` for
                 available modes.
             pad_constant_values (int, tuple(int)):
-                Constant values to use when padding with constant pad mode.
+                Constant values to use when padding with cmonstant pad mode.
                 Passed to the :func:`numpy.pad` `constant_values` argument.
                 Default is 0.
             coord_space (str):