Upload LoadImage.py

LoadImage.py

@@ -0,0 +1,161 @@
+import numpy as np
+import tifffile as tif
+import skimage.io as io
+from typing import Optional, Sequence, Union
+from monai.config import DtypeLike, PathLike, KeysCollection
+from monai.utils import ensure_tuple
+from monai.data.utils import is_supported_format, optional_import, ensure_tuple_rep
+from monai.data.image_reader import ImageReader, NumpyReader
+from monai.transforms import LoadImage, LoadImaged
+from monai.utils.enums import PostFix
+
+DEFAULT_POST_FIX = PostFix.meta()
+itk, has_itk = optional_import("itk", allow_namespace_pkg=True)
+
+__all__ = [
+    "CustomLoadImaged",
+    "CustomLoadImageD",
+    "CustomLoadImageDict",
+    "CustomLoadImage",
+]
+
+
+class CustomLoadImage(LoadImage):
+    """
+    Load image file or files from provided path based on reader.
+    If reader is not specified, this class automatically chooses readers
+    based on the supported suffixes and in the following order:
+
+        - User-specified reader at runtime when calling this loader.
+        - User-specified reader in the constructor of `LoadImage`.
+        - Readers from the last to the first in the registered list.
+        - Current default readers: (nii, nii.gz -> NibabelReader), (png, jpg, bmp -> PILReader),
+          (npz, npy -> NumpyReader), (nrrd -> NrrdReader), (DICOM file -> ITKReader).
+
+    [!Caution] This overriding replaces the original ITK with Custom UnifiedITKReader.
+    """
+
+    def __init__(
+        self,
+        reader=None,
+        image_only: bool = False,
+        dtype: DtypeLike = np.float32,
+        ensure_channel_first: bool = False,
+        *args,
+        **kwargs,
+    ) -> None:
+        super(CustomLoadImage, self).__init__(
+            reader, image_only, dtype, ensure_channel_first, *args, **kwargs
+        )
+
+        # Adding TIFFReader. Although ITK Reader supports ".tiff" files, sometimes fails to load images.
+        self.readers = []
+        self.register(UnifiedITKReader(*args, **kwargs))
+
+
+class CustomLoadImaged(LoadImaged):
+    """
+    Dictionary-based wrapper of `CustomLoadImage`.
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        reader: Optional[Union[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=False,
+        allow_missing_keys: bool = False,
+        *args,
+        **kwargs,
+    ) -> None:
+        super(CustomLoadImaged, self).__init__(
+            keys,
+            reader,
+            dtype,
+            meta_keys,
+            meta_key_postfix,
+            overwriting,
+            image_only,
+            ensure_channel_first,
+            simple_keys,
+            allow_missing_keys,
+            *args,
+            **kwargs,
+        )
+
+        # Assign CustomLoader
+        self._loader = CustomLoadImage(
+            reader, image_only, dtype, ensure_channel_first, *args, **kwargs
+        )
+        if not isinstance(meta_key_postfix, str):
+            raise TypeError(
+                f"meta_key_postfix must be a str but is {type(meta_key_postfix).__name__}."
+            )
+        self.meta_keys = (
+            ensure_tuple_rep(None, len(self.keys))
+            if meta_keys is None
+            else ensure_tuple(meta_keys)
+        )
+        if len(self.keys) != len(self.meta_keys):
+            raise ValueError("meta_keys should have the same length as keys.")
+        self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))
+        self.overwriting = overwriting
+
+
+class UnifiedITKReader(NumpyReader):
+    """
+    Unified Reader to read ".tif" and ".tiff files".
+    As the tifffile reads the images as numpy arrays, it inherits from the NumpyReader.
+    """
+
+    def __init__(
+        self, channel_dim: Optional[int] = None, **kwargs,
+    ):
+        super(UnifiedITKReader, self).__init__(channel_dim=channel_dim, **kwargs)
+        self.kwargs = kwargs
+        self.channel_dim = channel_dim
+
+    def verify_suffix(self, filename: Union[Sequence[PathLike], PathLike]) -> bool:
+        """Verify whether the file format is supported by TIFF Reader."""
+
+        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 Images from the file."""
+        img_ = []
+
+        filenames: Sequence[PathLike] = ensure_tuple(data)
+        kwargs_ = self.kwargs.copy()
+        kwargs_.update(kwargs)
+
+        for name in filenames:
+            name = f"{name}"
+
+            if name.endswith(".tif") or name.endswith(".tiff"):
+                _obj = tif.imread(name)
+            else:
+                try:
+                    _obj = itk.imread(name, **kwargs_)
+                    _obj = itk.array_view_from_image(_obj, keep_axes=False)
+                except:
+                    _obj = io.imread(name)
+
+            if len(_obj.shape) == 2:
+                _obj = np.repeat(np.expand_dims(_obj, axis=-1), 3, axis=-1)
+            elif len(_obj.shape) == 3 and _obj.shape[-1] > 3:
+                _obj = _obj[:, :, :3]
+            else:
+                pass
+
+            img_.append(_obj)
+
+        return img_ if len(filenames) > 1 else img_[0]
+
+
+CustomLoadImageD = CustomLoadImageDict = CustomLoadImaged