Download preprocessing.py from AI4MIRO/SwinUNETR_body_segmentation: direct link, hf CLI and curl.
- Browser
- Download file 15.9 kB
-
https://huggingface.co/AI4MIRO/SwinUNETR_body_segmentation/resolve/main/preprocessing.py
- Command line
-
hf download hf://AI4MIRO/SwinUNETR_body_segmentation/preprocessing.py
-
curl -L -o preprocessing.py https://huggingface.co/AI4MIRO/SwinUNETR_body_segmentation/resolve/main/preprocessing.py
15.9 kB
| from monai.transforms import MapTransform, Transform | |
| import numpy as np | |
| from monai.config import KeysCollection | |
| from typing import Dict, Hashable, Mapping, Optional, Type, Union, Sequence | |
| import sys | |
| from pathlib import Path | |
| from monai.config import DtypeLike, KeysCollection | |
| from monai.data.image_reader import ImageReader | |
| from monai.transforms.transform import MapTransform | |
| from monai.utils import ensure_tuple, ensure_tuple_rep | |
| from monai.utils.enums import PostFix | |
| from monai.data.meta_tensor import MetaTensor | |
| from monai.transforms.utility.array import EnsureChannelFirst | |
| from monai.utils import ImageMetaKey as Key | |
| from monai.utils import ensure_tuple, ensure_tuple_rep, convert_to_dst_type | |
| from monai.data import NibabelReader | |
| from monai.config import DtypeLike, KeysCollection, PathLike, NdarrayOrTensor | |
| DEFAULT_POST_FIX = PostFix.meta() | |
| class LoadImaged(MapTransform): | |
| """ | |
| Dictionary-based wrapper of :py:class:`monai.transforms.LoadImage`, | |
| It can load both image data and metadata. When loading a list of files in one key, | |
| the arrays will be stacked and a new dimension will be added as the first dimension | |
| In this case, the metadata of the first image will be used to represent the stacked result. | |
| The affine transform of all the stacked images should be same. | |
| The output metadata field will be created as ``meta_keys`` or ``key_{meta_key_postfix}``. | |
| 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), (dcm, DICOM series and others -> ITKReader). | |
| Please note that for png, jpg, bmp, and other 2D formats, readers often swap axis 0 and 1 after | |
| loading the array because the `HW` definition for non-medical specific file formats is different | |
| from other common medical packages. | |
| Note: | |
| - If `reader` is specified, the loader will attempt to use the specified readers and the default supported | |
| readers. This might introduce overheads when handling the exceptions of trying the incompatible loaders. | |
| In this case, it is therefore recommended setting the most appropriate reader as | |
| the last item of the `reader` parameter. | |
| See also: | |
| - tutorial: https://github.com/Project-MONAI/tutorials/blob/master/modules/load_medical_images.ipynb | |
| """ | |
| def __init__( | |
| self, | |
| keys: KeysCollection, | |
| reader: Optional[Union[ImageReader, str]] = None, | |
| patientname: str='', | |
| 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, | |
| prune_meta_pattern: Optional[str] = None, | |
| prune_meta_sep: str = ".", | |
| allow_missing_keys: bool = False, | |
| *args, | |
| **kwargs, | |
| ) -> None: | |
| """ | |
| Args: | |
| keys: keys of the corresponding items to be transformed. | |
| See also: :py:class:`monai.transforms.compose.MapTransform` | |
| reader: reader to load image file and metadata | |
| - if `reader` is None, a default set of `SUPPORTED_READERS` will be used. | |
| - if `reader` is a string, it's treated as a class name or dotted path | |
| (such as ``"monai.data.ITKReader"``), the supported built-in reader classes are | |
| ``"ITKReader"``, ``"NibabelReader"``, ``"NumpyReader"``. | |
| a reader instance will be constructed with the `*args` and `**kwargs` parameters. | |
| - if `reader` is a reader class/instance, it will be registered to this loader accordingly. | |
| patientname: the patient name. | |
| dtype: if not None, convert the loaded image data to this data type. | |
| meta_keys: explicitly indicate the key to store the corresponding metadata dictionary. | |
| the metadata is a dictionary object which contains: filename, original_shape, etc. | |
| it can be a sequence of string, map to the `keys`. | |
| if None, will try to construct meta_keys by `key_{meta_key_postfix}`. | |
| meta_key_postfix: if meta_keys is None, use `key_{postfix}` to store the metadata of the nifti image, | |
| default is `meta_dict`. The metadata is a dictionary object. | |
| For example, load nifti file for `image`, store the metadata into `image_meta_dict`. | |
| overwriting: whether allow overwriting existing metadata of same key. | |
| default is False, which will raise exception if encountering existing key. | |
| image_only: if True return dictionary containing just only the image volumes, otherwise return | |
| dictionary containing image data array and header dict per input key. | |
| ensure_channel_first: if `True` and loaded both image array and metadata, automatically convert | |
| the image array shape to `channel first`. default to `False`. | |
| simple_keys: whether to remove redundant metadata keys, default to False for backward compatibility. | |
| prune_meta_pattern: combined with `prune_meta_sep`, a regular expression used to match and prune keys | |
| in the metadata (nested dictionary), default to None, no key deletion. | |
| prune_meta_sep: combined with `prune_meta_pattern`, used to match and prune keys | |
| in the metadata (nested dictionary). default is ".", see also :py:class:`monai.transforms.DeleteItemsd`. | |
| e.g. ``prune_meta_pattern=".*_code$", prune_meta_sep=" "`` removes meta keys that ends with ``"_code"``. | |
| allow_missing_keys: don't raise exception if key is missing. | |
| args: additional parameters for reader if providing a reader name. | |
| kwargs: additional parameters for reader if providing a reader name. | |
| """ | |
| super().__init__(keys, allow_missing_keys) | |
| self._loader = LoadImage( | |
| reader, | |
| patientname, | |
| image_only, | |
| dtype, | |
| ensure_channel_first, | |
| simple_keys, | |
| prune_meta_pattern, | |
| prune_meta_sep, | |
| *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 | |
| # Check patient name | |
| if (len(patientname)==0): | |
| raise ValueError("Patient name should not be empty.") | |
| def register(self, reader: ImageReader): | |
| self._loader.register(reader) | |
| def __call__(self, data, reader: Optional[ImageReader] = None): | |
| """ | |
| Raises: | |
| KeyError: When not ``self.overwriting`` and key already exists in ``data``. | |
| """ | |
| d = dict(data) | |
| for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): | |
| data = self._loader(d[key], reader) | |
| if self._loader.image_only: | |
| d[key] = data | |
| else: | |
| if not isinstance(data, (tuple, list)): | |
| raise ValueError("loader must return a tuple or list (because image_only=False was used).") | |
| d[key] = data[0] | |
| if not isinstance(data[1], dict): | |
| raise ValueError("metadata must be a dict.") | |
| meta_key = meta_key or f"{key}_{meta_key_postfix}" | |
| if meta_key in d and not self.overwriting: | |
| raise KeyError(f"Metadata with key {meta_key} already exists and overwriting=False.") | |
| d[meta_key] = data[1] | |
| return d | |
| def switch_endianness(data, new="<"): | |
| """ | |
| Convert the input `data` endianness to `new`. | |
| Args: | |
| data: input to be converted. | |
| new: the target endianness, currently support "<" or ">". | |
| """ | |
| if isinstance(data, np.ndarray): | |
| # default to system endian | |
| sys_native = "<" if (sys.byteorder == "little") else ">" | |
| current_ = sys_native if data.dtype.byteorder not in ("<", ">") else data.dtype.byteorder | |
| if new not in ("<", ">"): | |
| raise NotImplementedError(f"Not implemented option new={new}.") | |
| if current_ != new: | |
| data = data.byteswap().newbyteorder(new) | |
| elif isinstance(data, tuple): | |
| data = tuple(switch_endianness(x, new) for x in data) | |
| elif isinstance(data, list): | |
| data = [switch_endianness(x, new) for x in data] | |
| elif isinstance(data, dict): | |
| data = {k: switch_endianness(v, new) for k, v in data.items()} | |
| elif not isinstance(data, (bool, str, float, int, type(None))): | |
| raise RuntimeError(f"Unknown type: {type(data).__name__}") | |
| return data | |
| class LoadImage(Transform): | |
| """ | |
| 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). | |
| Please note that for png, jpg, bmp, and other 2D formats, readers often swap axis 0 and 1 after | |
| loading the array because the `HW` definition for non-medical specific file formats is different | |
| from other common medical packages. | |
| See also: | |
| - tutorial: https://github.com/Project-MONAI/tutorials/blob/master/modules/load_medical_images.ipynb | |
| """ | |
| def __init__( | |
| self, | |
| reader="NibabelReader", | |
| patientname:str="", | |
| image_only: bool = False, | |
| dtype: DtypeLike = np.float32, | |
| ensure_channel_first: bool = False, | |
| simple_keys: bool = False, | |
| prune_meta_pattern: Optional[str] = None, | |
| prune_meta_sep: str = ".", | |
| *args, | |
| **kwargs, | |
| ) -> None: | |
| """ | |
| Args: | |
| reader: reader to load image file and metadata | |
| - if `reader` is None, a default set of `SUPPORTED_READERS` will be used. | |
| - if `reader` is a string, it's treated as a class name or dotted path | |
| (such as ``"monai.data.ITKReader"``), the supported built-in reader classes are | |
| ``"ITKReader"``, ``"NibabelReader"``, ``"NumpyReader"``, ``"PydicomReader"``. | |
| a reader instance will be constructed with the `*args` and `**kwargs` parameters. | |
| - if `reader` is a reader class/instance, it will be registered to this loader accordingly. | |
| image_only: if True return only the image MetaTensor, otherwise return image and header dict. | |
| dtype: if not None convert the loaded image to this data type. | |
| ensure_channel_first: if `True` and loaded both image array and metadata, automatically convert | |
| the image array shape to `channel first`. default to `False`. | |
| simple_keys: whether to remove redundant metadata keys, default to False for backward compatibility. | |
| prune_meta_pattern: combined with `prune_meta_sep`, a regular expression used to match and prune keys | |
| in the metadata (nested dictionary), default to None, no key deletion. | |
| prune_meta_sep: combined with `prune_meta_pattern`, used to match and prune keys | |
| in the metadata (nested dictionary). default is ".", see also :py:class:`monai.transforms.DeleteItemsd`. | |
| e.g. ``prune_meta_pattern=".*_code$", prune_meta_sep=" "`` removes meta keys that ends with ``"_code"``. | |
| args: additional parameters for reader if providing a reader name. | |
| kwargs: additional parameters for reader if providing a reader name. | |
| Note: | |
| - The transform returns a MetaTensor, unless `set_track_meta(False)` has been used, in which case, a | |
| `torch.Tensor` will be returned. | |
| - If `reader` is specified, the loader will attempt to use the specified readers and the default supported | |
| readers. This might introduce overheads when handling the exceptions of trying the incompatible loaders. | |
| In this case, it is therefore recommended setting the most appropriate reader as | |
| the last item of the `reader` parameter. | |
| """ | |
| self.auto_select = reader is None | |
| self.image_only = image_only | |
| self.dtype = dtype | |
| self.ensure_channel_first = ensure_channel_first | |
| self.simple_keys = simple_keys | |
| self.pattern = prune_meta_pattern | |
| self.sep = prune_meta_sep | |
| self.patientname = patientname | |
| return | |
| def __call__(self, filename: Union[Sequence[PathLike], PathLike], reader: Optional[ImageReader] = NibabelReader): | |
| """ | |
| Load image file and metadata from the given filename(s). | |
| If `reader` is not specified, this class automatically chooses readers based on the | |
| reversed order of registered readers `self.readers`. | |
| Args: | |
| filename: path file or file-like object or a list of files. | |
| will save the filename to meta_data with key `filename_or_obj`. | |
| if provided a list of files, use the filename of first file to save, | |
| and will stack them together as multi-channels data. | |
| if provided directory path instead of file path, will treat it as | |
| DICOM images series and read. | |
| reader: runtime reader to load image file and metadata. | |
| """ | |
| filename = tuple(f"{Path(s).expanduser()}" for s in ensure_tuple(filename)) # allow Path objects | |
| img, err = None, [] | |
| # if reader is not None: | |
| reader=NibabelReader() | |
| img = reader.read(filename) # runtime specified reader | |
| header = img.header | |
| img_array = img.get_fdata() | |
| img_array: NdarrayOrTensor | |
| img_array, meta_data = reader.get_data(img) | |
| img_array = convert_to_dst_type(img_array, dst=img_array, dtype=self.dtype)[0] | |
| if not isinstance(meta_data, dict): | |
| raise ValueError("`meta_data` must be a dict.") | |
| # make sure all elements in metadata are little endian | |
| meta_data = switch_endianness(meta_data, "<") | |
| # meta_data["patient_name"]=filename[0].split('/')[-3] | |
| meta_data["patient_name"]=self.patientname | |
| meta_data[Key.FILENAME_OR_OBJ] = f"{ensure_tuple(filename)[0]}" # Path obj should be strings for data loader | |
| img = MetaTensor.ensure_torch_and_prune_meta( | |
| img_array, meta_data, self.simple_keys, pattern=self.pattern, sep=self.sep | |
| ) | |
| if self.ensure_channel_first: | |
| img = EnsureChannelFirst()(img) | |
| if self.image_only: | |
| return img | |
| return img, img.meta, header if isinstance(img, MetaTensor) else meta_data |