import numpy as np import tifffile as tif import skimage.io as io from typing import List, Optional, Sequence, Type, Union from monai.utils.enums import PostFix from monai.utils.module import optional_import from monai.utils.misc import ensure_tuple, ensure_tuple_rep from monai.data.utils import is_supported_format from monai.data.image_reader import ImageReader, NumpyReader from monai.transforms import LoadImage, LoadImaged # type: ignore from monai.config.type_definitions import DtypeLike, PathLike, KeysCollection # Default value for metadata postfix DEFAULT_POST_FIX = PostFix.meta() # Try to import ITK library; if not available, has_itk will be False itk, has_itk = optional_import("itk", allow_namespace_pkg=True) __all__ = [ "CustomLoadImage", # Basic image loader "CustomLoadImaged", # Dictionary-based image loader "CustomLoadImageD", # Dictionary-based image loader "CustomLoadImageDict", # Dictionary-based image loader ] class CustomLoadImage(LoadImage): """ Class for loading one or multiple images from a given path. If a reader is not specified, the appropriate file reading method is automatically chosen based on the file extension. Priority: - Reader passed by the user at runtime. - Reader specified in the constructor. - Registered readers (from last to first). - Standard readers for different formats (e.g., NibabelReader for nii, PILReader for png/jpg, etc.). [Note] Here, the original ITKReader is replaced by the universal reader UniversalImageReader. """ def __init__( self, reader: Optional[Union[ImageReader, Type[ImageReader], str]] = None, image_only: bool = False, dtype: DtypeLike = np.float32, ensure_channel_first: bool = False, *args, **kwargs, ) -> None: super().__init__( reader=reader, image_only=image_only, dtype=dtype, ensure_channel_first=ensure_channel_first, *args, **kwargs ) # Clear the list of registered readers self.readers = [] # Register the universal reader that handles TIFF, PNG, JPG, BMP, etc. self.register(UniversalImageReader(*args, **kwargs)) class CustomLoadImaged(LoadImaged): """ Dictionary-based image loader. Wraps image loading with CustomLoadImage and allows processing of data represented as a dictionary, where keys point to file paths. """ def __init__( self, keys: KeysCollection, reader: Optional[Union[Type[ImageReader], str]] = None, dtype: DtypeLike = np.float32, meta_keys: Optional[KeysCollection] = None, meta_key_postfix: str = DEFAULT_POST_FIX, overwriting: bool = False, image_only: bool = False, ensure_channel_first: bool = False, simple_keys: bool = False, allow_missing_keys: bool = False, *args, **kwargs, ) -> None: super().__init__( keys=keys, reader=reader, dtype=dtype, meta_keys=meta_keys, meta_key_postfix=meta_key_postfix, overwriting=overwriting, image_only=image_only, ensure_channel_first=ensure_channel_first, simple_keys=simple_keys, allow_missing_keys=allow_missing_keys, *args, **kwargs, ) # Assign the custom image loader self._loader = CustomLoadImage( reader=reader, image_only=image_only, dtype=dtype, ensure_channel_first=ensure_channel_first, *args, **kwargs ) # Ensure that meta_key_postfix is a string if not isinstance(meta_key_postfix, str): raise TypeError( f"meta_key_postfix must be a string, but got {type(meta_key_postfix).__name__}." ) # If meta_keys are not provided, create a tuple of None for each key self.meta_keys = ( ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys) ) # Check that the number of meta_keys matches the number of keys if len(self.keys) != len(self.meta_keys): raise ValueError("meta_keys must have the same length as keys.") # Assign each key its corresponding metadata postfix self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) self.overwriting = overwriting class UniversalImageReader(NumpyReader): """ Universal image reader for TIFF, PNG, JPG, BMP, etc. Uses: - tifffile for reading TIFF files. - ITK (if available) for reading other formats. - skimage.io for reading if the previous methods fail. The image is loaded with its original number of channels (layers) without forced modifications (e.g., repeating or cropping channels). """ def __init__( self, channel_dim: Optional[int] = None, **kwargs, ): super().__init__(channel_dim=channel_dim, **kwargs) self.kwargs = kwargs self.channel_dim = channel_dim def verify_suffix(self, filename: Union[Sequence[PathLike], PathLike]) -> bool: """ Check if the file format is supported for reading. Supported extensions: tif, tiff, png, jpg, bmp, jpeg. """ suffixes: Sequence[str] = ["tif", "tiff", "png", "jpg", "bmp", "jpeg"] return has_itk or is_supported_format(filename, suffixes) def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs): """ Read image(s) from the given path. Arguments: data: A file path or a sequence of file paths. kwargs: Additional parameters for reading. Returns: A single image or a list of images depending on the number of paths provided. """ images: List[np.ndarray] = [] # List to store the loaded images # Convert data to a tuple to support multiple files filenames: Sequence[PathLike] = ensure_tuple(data) # Merge parameters provided in the constructor and the read() method kwargs_ = self.kwargs.copy() kwargs_.update(kwargs) for name in filenames: # Convert file name to string name = f"{name}" # If the file has a .tif or .tiff extension (case-insensitive), use tifffile for reading if name.lower().endswith((".tif", ".tiff")): img_array = tif.imread(name) else: # Attempt to read the image using ITK (if available) try: img_itk = itk.imread(name, **kwargs_) img_array = itk.array_view_from_image(img_itk, keep_axes=False) except Exception: # If ITK fails, use skimage.io for reading img_array = io.imread(name) # Check the number of dimensions (axes) of the loaded image if img_array.ndim == 2: # If the image is 2D (height, width), add a new axis at the end to represent the channel img_array = np.expand_dims(img_array, axis=-1) images.append(img_array) # Return a single image if only one file was provided, otherwise return a list of images return images if len(filenames) > 1 else images[0] CustomLoadImageD = CustomLoadImageDict = CustomLoadImaged