Spaces:
Running
Running
File size: 5,431 Bytes
a4b3c40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
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
|