Improve CuPy type hinting in ImageReader

Signed-off-by: ytl0623 <david89062388@gmail.com>

Author: ytl0623

monai/data/image_reader.py

@@ -22,7 +22,7 @@
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Union
 
 import numpy as np
 from torch.utils.data._utils.collate import np_str_obj_array_pattern
@@ -38,14 +38,17 @@
 from monai.utils import MetaKeys, SpaceKeys, TraceKeys, ensure_tuple, optional_import, require_pkg
 
 if TYPE_CHECKING:
+    import cupy as cp
     import itk
     import nibabel as nib
     import nrrd
     import pydicom
     from nibabel.nifti1 import Nifti1Image
     from PIL import Image as PILImage
 
-    has_nrrd = has_itk = has_nib = has_pil = has_pydicom = True
+    has_nrrd = has_itk = has_nib = has_pil = has_pydicom, has_cp = True
+    Ndarray = Union[np.ndarray, cp.ndarray]
+
 else:
     itk, has_itk = optional_import("itk", allow_namespace_pkg=True)
     nib, has_nib = optional_import("nibabel")
@@ -54,7 +57,12 @@
     pydicom, has_pydicom = optional_import("pydicom")
     nrrd, has_nrrd = optional_import("nrrd", allow_namespace_pkg=True)
 
-cp, has_cp = optional_import("cupy")
+    cp, has_cp = optional_import("cupy")
+    if has_cp:
+        Ndarray = Union[np.ndarray, cp.ndarray]
+    else:
+        Ndarray = np.ndarray
+
 kvikio, has_kvikio = optional_import("kvikio")
 
 __all__ = ["ImageReader", "ITKReader", "NibabelReader", "NumpyReader", "PILReader", "PydicomReader", "NrrdReader"]
@@ -107,7 +115,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs) -> Sequence[Any] |
         raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")
 
     @abstractmethod
-    def get_data(self, img) -> tuple[np.ndarray, dict]:
+    def get_data(self, img) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function must return two objects, the first is a numpy array of image data,
@@ -143,7 +151,7 @@ def _copy_compatible_dict(from_dict: dict, to_dict: dict):
             )
 
 
-def _stack_images(image_list: list, meta_dict: dict, to_cupy: bool = False):
+def _stack_images(image_list: list[Ndarray], meta_dict: dict, to_cupy: bool = False) -> Ndarray:
     if len(image_list) <= 1:
         return image_list[0]
     if not is_no_channel(meta_dict.get(MetaKeys.ORIGINAL_CHANNEL_DIM, None)):
@@ -269,7 +277,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs):
                 img_.append(itk.imread(name, **kwargs_))
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img) -> tuple[np.ndarray, dict]:
+    def get_data(self, img) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -281,7 +289,7 @@ def get_data(self, img) -> tuple[np.ndarray, dict]:
             img: an ITK image object loaded from an image file or a list of ITK image objects.
 
         """
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
 
         for i in ensure_tuple(img):
@@ -616,7 +624,7 @@ def _combine_dicom_series(self, data: Iterable, filenames: Sequence[PathLike]):
 
         return stack_array, stack_metadata
 
-    def get_data(self, data) -> tuple[np.ndarray, dict]:
+    def get_data(self, data) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -663,10 +671,7 @@ def get_data(self, data) -> tuple[np.ndarray, dict]:
                     metadata[MetaKeys.SPATIAL_SHAPE] = data_array.shape
                 dicom_data.append((data_array, metadata))
 
-        # TODO: the actual type is list[np.ndarray | cp.ndarray]
-        # should figure out how to define correct types without having cupy not found error
-        # https://github.com/Project-MONAI/MONAI/pull/8188#discussion_r1886645918
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
 
         for data_array, metadata in ensure_tuple(dicom_data):
@@ -841,7 +846,7 @@ def _get_seg_data(self, img, filename):
         if self.label_dict is not None:
             metadata["labels"] = self.label_dict
             if self.to_gpu:
-                all_segs = cp.zeros([*spatial_shape, len(self.label_dict)], dtype=array_data.dtype)
+                all_segs: Ndarray = cp.zeros([*spatial_shape, len(self.label_dict)], dtype=array_data.dtype)
             else:
                 all_segs = np.zeros([*spatial_shape, len(self.label_dict)], dtype=array_data.dtype)
         else:
@@ -899,7 +904,7 @@ def _get_seg_data(self, img, filename):
 
         return all_segs, metadata
 
-    def _get_array_data_from_gpu(self, img, filename):
+    def _get_array_data_from_gpu(self, img, filename) -> Ndarray:
         """
         Get the raw array data of the image. This function is used when `to_gpu` is set to True.
 
@@ -954,7 +959,7 @@ def _get_array_data_from_gpu(self, img, filename):
 
         return data
 
-    def _get_array_data(self, img, filename):
+    def _get_array_data(self, img, filename) -> Ndarray:
         """
         Get the array data of the image. If `RescaleSlope` and `RescaleIntercept` are available, the raw array data
         will be rescaled. The output data has the dtype float32 if the rescaling is applied.
@@ -1092,7 +1097,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs):
             img_.append(img)  # type: ignore
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img) -> tuple[np.ndarray, dict]:
+    def get_data(self, img) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -1104,10 +1109,7 @@ def get_data(self, img) -> tuple[np.ndarray, dict]:
             img: a Nibabel image object loaded from an image file or a list of Nibabel image objects.
 
         """
-        # TODO: the actual type is list[np.ndarray | cp.ndarray]
-        # should figure out how to define correct types without having cupy not found error
-        # https://github.com/Project-MONAI/MONAI/pull/8188#discussion_r1886645918
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
 
         for i, filename in zip(ensure_tuple(img), self.filenames):
@@ -1186,7 +1188,7 @@ def _get_spatial_shape(self, img):
         spatial_rank = max(min(ndim, 3), 1)
         return np.asarray(size[:spatial_rank])
 
-    def _get_array_data(self, img, filename):
+    def _get_array_data(self, img, filename) -> Ndarray:
         """
         Get the raw array data of the image, converted to Numpy array.
 
@@ -1281,7 +1283,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs):
 
         return img_ if len(img_) > 1 else img_[0]
 
-    def get_data(self, img) -> tuple[np.ndarray, dict]:
+    def get_data(self, img) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -1293,7 +1295,7 @@ def get_data(self, img) -> tuple[np.ndarray, dict]:
             img: a Numpy array loaded from a file or a list of Numpy arrays.
 
         """
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
         if isinstance(img, np.ndarray):
             img = (img,)
@@ -1374,7 +1376,7 @@ def read(self, data: Sequence[PathLike] | PathLike | np.ndarray, **kwargs):
 
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img) -> tuple[np.ndarray, dict]:
+    def get_data(self, img) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -1388,7 +1390,7 @@ def get_data(self, img) -> tuple[np.ndarray, dict]:
             img: a PIL Image object loaded from a file or a list of PIL Image objects.
 
         """
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
 
         for i in ensure_tuple(img):
@@ -1425,7 +1427,7 @@ def _get_spatial_shape(self, img):
 class NrrdImage:
     """Class to wrap nrrd image array and metadata header"""
 
-    array: np.ndarray
+    array: Ndarray
     header: dict
 
 
@@ -1495,7 +1497,7 @@ def read(self, data: Sequence[PathLike] | PathLike, **kwargs) -> Sequence[Any] |
             img_.append(nrrd_image)
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img: NrrdImage | list[NrrdImage]) -> tuple[np.ndarray, dict]:
+    def get_data(self, img: NrrdImage | list[NrrdImage]) -> tuple[Ndarray, dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function must return two objects, the first is a numpy array of image data,
@@ -1505,7 +1507,7 @@ def get_data(self, img: NrrdImage | list[NrrdImage]) -> tuple[np.ndarray, dict]:
             img: a `NrrdImage` loaded from an image file or a list of image objects.
 
         """
-        img_array: list[np.ndarray] = []
+        img_array: list[Ndarray] = []
         compatible_meta: dict = {}
 
         for i in ensure_tuple(img):