Download postprocessing.py from AI4MIRO/SwinUNETR_body_segmentation: direct link, hf CLI and curl.
- Browser
- Download file 18.1 kB
-
https://huggingface.co/AI4MIRO/SwinUNETR_body_segmentation/resolve/main/postprocessing.py
- Command line
-
hf download hf://AI4MIRO/SwinUNETR_body_segmentation/postprocessing.py
-
curl -L -o postprocessing.py https://huggingface.co/AI4MIRO/SwinUNETR_body_segmentation/resolve/main/postprocessing.py
18.1 kB
| from monai.transforms import MapTransform, Transform | |
| from monai.config import KeysCollection | |
| from typing import Dict, Hashable, Mapping, Optional, Type, Union, Sequence | |
| import torch, sys | |
| from pathlib import Path | |
| from monai.config import DtypeLike, KeysCollection, PathLike | |
| from monai.data import image_writer | |
| from monai.transforms.transform import MapTransform | |
| from monai.utils import GridSamplePadMode, ensure_tuple, ensure_tuple_rep, optional_import | |
| from monai.data.meta_tensor import MetaTensor | |
| from monai.data.folder_layout import FolderLayout | |
| from pydoc import locate | |
| import numpy as np | |
| import nibabel as nib, os | |
| from monai.utils.enums import PostFix | |
| DEFAULT_POST_FIX = PostFix.meta() | |
| def set_header_info(nii_file, voxelsize, image_position_patient, contours_exist = None): | |
| nii_file.header['pixdim'][1] = voxelsize[0] | |
| nii_file.header['pixdim'][2] = voxelsize[1] | |
| nii_file.header['pixdim'][3] = voxelsize[2] | |
| #affine - voxelsize | |
| nii_file.affine[0][0] = voxelsize[0] | |
| nii_file.affine[1][1] = voxelsize[1] | |
| nii_file.affine[2][2] = voxelsize[2] | |
| #affine - imagecorner | |
| nii_file.affine[0][3] = image_position_patient[0] | |
| nii_file.affine[1][3] = image_position_patient[1] | |
| nii_file.affine[2][3] = image_position_patient[2] | |
| if contours_exist: | |
| nii_file.header.extensions.append(nib.nifti1.Nifti1Extension(0, bytearray(contours_exist))) | |
| return nii_file | |
| def add_contours_exist(preddir, refCT): | |
| img = nib.load(os.path.join(preddir, 'RTStruct.nii.gz')) | |
| data = img.get_fdata().astype(int) | |
| contours_exist = [] | |
| data_one_hot = np.zeros(data.shape[:3]) | |
| # We remove the first channel as it is the background | |
| for i in range(data.shape[-1]-1): | |
| if np.count_nonzero(data[:,:,:,i+1])>0: | |
| contours_exist.append(1) | |
| data_one_hot+=np.where(data[:,:,:,i+1]==1,2**i,0) | |
| else: | |
| contours_exist.append(0) | |
| data_one_hot_nii = nib.Nifti1Image(data_one_hot, affine=np.eye(4)) | |
| data_one_hot_nii = set_header_info(data_one_hot_nii, voxelsize=np.array(refCT.PixelSpacing), image_position_patient=refCT.ImagePositionPatient, contours_exist=contours_exist) | |
| nib.save(data_one_hot_nii,os.path.join(preddir, 'RTStruct.nii.gz')) | |
| class SaveImaged(MapTransform): | |
| """ | |
| Dictionary-based wrapper of :py:class:`monai.transforms.SaveImage`. | |
| Note: | |
| Image should be channel-first shape: [C,H,W,[D]]. | |
| If the data is a patch of big image, will append the patch index to filename. | |
| Args: | |
| keys: keys of the corresponding items to be transformed. | |
| See also: :py:class:`monai.transforms.compose.MapTransform` | |
| meta_keys: explicitly indicate the key of the corresponding metadata dictionary. | |
| For example, for data with key `image`, the metadata by default is in `image_meta_dict`. | |
| The metadata is a dictionary contains values such as filename, original_shape. | |
| This argument 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_{meta_key_postfix}` to retrieve the metadict. | |
| output_dir: output image directory. | |
| output_postfix: a string appended to all output file names, default to `trans`. | |
| output_ext: output file extension name, available extensions: `.nii.gz`, `.nii`, `.png`. | |
| output_dtype: data type for saving data. Defaults to ``np.float32``. | |
| resample: whether to resample image (if needed) before saving the data array, | |
| based on the `spatial_shape` (and `original_affine`) from metadata. | |
| mode: This option is used when ``resample=True``. Defaults to ``"nearest"``. | |
| Depending on the writers, the possible options are: | |
| - {``"bilinear"``, ``"nearest"``, ``"bicubic"``}. | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample | |
| - {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}. | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#interpolate | |
| padding_mode: This option is used when ``resample = True``. Defaults to ``"border"``. | |
| Possible options are {``"zeros"``, ``"border"``, ``"reflection"``} | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample | |
| scale: {``255``, ``65535``} postprocess data by clipping to [0, 1] and scaling | |
| [0, 255] (uint8) or [0, 65535] (uint16). Default is `None` (no scaling). | |
| dtype: data type during resampling computation. Defaults to ``np.float64`` for best precision. | |
| if None, use the data type of input data. To be compatible with other modules, | |
| output_dtype: data type for saving data. Defaults to ``np.float32``. | |
| it's used for NIfTI format only. | |
| allow_missing_keys: don't raise exception if key is missing. | |
| squeeze_end_dims: if True, any trailing singleton dimensions will be removed (after the channel | |
| has been moved to the end). So if input is (C,H,W,D), this will be altered to (H,W,D,C), and | |
| then if C==1, it will be saved as (H,W,D). If D is also 1, it will be saved as (H,W). If `false`, | |
| image will always be saved as (H,W,D,C). | |
| data_root_dir: if not empty, it specifies the beginning parts of the input file's | |
| absolute path. It's used to compute `input_file_rel_path`, the relative path to the file from | |
| `data_root_dir` to preserve folder structure when saving in case there are files in different | |
| folders with the same file names. For example, with the following inputs: | |
| - input_file_name: `/foo/bar/test1/image.nii` | |
| - output_postfix: `seg` | |
| - output_ext: `.nii.gz` | |
| - output_dir: `/output` | |
| - data_root_dir: `/foo/bar` | |
| The output will be: /output/test1/image/image_seg.nii.gz | |
| separate_folder: whether to save every file in a separate folder. For example: for the input filename | |
| `image.nii`, postfix `seg` and folder_path `output`, if `separate_folder=True`, it will be saved as: | |
| `output/image/image_seg.nii`, if `False`, saving as `output/image_seg.nii`. Default to `True`. | |
| print_log: whether to print logs when saving. Default to `True`. | |
| output_format: an optional string to specify the output image writer. | |
| see also: `monai.data.image_writer.SUPPORTED_WRITERS`. | |
| writer: a customised `monai.data.ImageWriter` subclass to save data arrays. | |
| if `None`, use the default writer from `monai.data.image_writer` according to `output_ext`. | |
| if it's a string, it's treated as a class name or dotted path; | |
| the supported built-in writer classes are ``"NibabelWriter"``, ``"ITKWriter"``, ``"PILWriter"``. | |
| """ | |
| def __init__( | |
| self, | |
| keys: KeysCollection, | |
| meta_keys: Optional[KeysCollection] = None, | |
| meta_key_postfix: str = DEFAULT_POST_FIX, | |
| output_dir: Union[Path, str] = "./", | |
| output_postfix: str = "trans", | |
| output_ext: str = ".nii.gz", | |
| resample: bool = True, | |
| mode: str = "nearest", | |
| padding_mode: str = GridSamplePadMode.BORDER, | |
| scale: Optional[int] = None, | |
| dtype: DtypeLike = np.float64, | |
| output_dtype: DtypeLike = np.float32, | |
| allow_missing_keys: bool = False, | |
| squeeze_end_dims: bool = True, | |
| data_root_dir: str = "", | |
| separate_folder: bool = True, | |
| print_log: bool = True, | |
| output_format: str = "", | |
| writer: Union[Type[image_writer.ImageWriter], str, None] = None, | |
| ) -> None: | |
| super().__init__(keys, allow_missing_keys) | |
| self.meta_keys = ensure_tuple_rep(meta_keys, len(self.keys)) | |
| self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys)) | |
| self.saver = SaveImage( | |
| output_dir=output_dir, | |
| output_postfix=output_postfix, | |
| output_ext=output_ext, | |
| resample=resample, | |
| mode=mode, | |
| padding_mode=padding_mode, | |
| scale=scale, | |
| dtype=dtype, | |
| output_dtype=output_dtype, | |
| squeeze_end_dims=squeeze_end_dims, | |
| data_root_dir=data_root_dir, | |
| separate_folder=separate_folder, | |
| print_log=print_log, | |
| output_format=output_format, | |
| writer=writer, | |
| ) | |
| def set_options(self, init_kwargs=None, data_kwargs=None, meta_kwargs=None, write_kwargs=None): | |
| self.saver.set_options(init_kwargs, data_kwargs, meta_kwargs, write_kwargs) | |
| def __call__(self, data): | |
| d = dict(data) | |
| for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix): | |
| if meta_key is None and meta_key_postfix is not None: | |
| meta_key = f"{key}_{meta_key_postfix}" | |
| meta_data = d.get(meta_key) if meta_key is not None else None | |
| self.saver(img=d[key], meta_data=meta_data) | |
| return d | |
| class SaveImage(Transform): | |
| """ | |
| Save the image (in the form of torch tensor or numpy ndarray) and metadata dictionary into files. | |
| The name of saved file will be `{input_image_name}_{output_postfix}{output_ext}`, | |
| where the `input_image_name` is extracted from the provided metadata dictionary. | |
| If no metadata provided, a running index starting from 0 will be used as the filename prefix. | |
| Args: | |
| output_dir: output image directory. | |
| output_postfix: a string appended to all output file names, default to `trans`. | |
| output_ext: output file extension name. | |
| output_dtype: data type for saving data. Defaults to ``np.float32``. | |
| resample: whether to resample image (if needed) before saving the data array, | |
| based on the `spatial_shape` (and `original_affine`) from metadata. | |
| mode: This option is used when ``resample=True``. Defaults to ``"nearest"``. | |
| Depending on the writers, the possible options are | |
| - {``"bilinear"``, ``"nearest"``, ``"bicubic"``}. | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample | |
| - {``"nearest"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}. | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#interpolate | |
| padding_mode: This option is used when ``resample = True``. Defaults to ``"border"``. | |
| Possible options are {``"zeros"``, ``"border"``, ``"reflection"``} | |
| See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample | |
| scale: {``255``, ``65535``} postprocess data by clipping to [0, 1] and scaling | |
| [0, 255] (uint8) or [0, 65535] (uint16). Default is `None` (no scaling). | |
| dtype: data type during resampling computation. Defaults to ``np.float64`` for best precision. | |
| if None, use the data type of input data. To be compatible with other modules, | |
| squeeze_end_dims: if True, any trailing singleton dimensions will be removed (after the channel | |
| has been moved to the end). So if input is (C,H,W,D), this will be altered to (H,W,D,C), and | |
| then if C==1, it will be saved as (H,W,D). If D is also 1, it will be saved as (H,W). If `false`, | |
| image will always be saved as (H,W,D,C). | |
| data_root_dir: if not empty, it specifies the beginning parts of the input file's | |
| absolute path. It's used to compute `input_file_rel_path`, the relative path to the file from | |
| `data_root_dir` to preserve folder structure when saving in case there are files in different | |
| folders with the same file names. For example, with the following inputs: | |
| - input_file_name: `/foo/bar/test1/image.nii` | |
| - output_postfix: `seg` | |
| - output_ext: `.nii.gz` | |
| - output_dir: `/output` | |
| - data_root_dir: `/foo/bar` | |
| The output will be: /output/test1/image/image_seg.nii.gz | |
| separate_folder: whether to save every file in a separate folder. For example: for the input filename | |
| `image.nii`, postfix `seg` and folder_path `output`, if `separate_folder=True`, it will be saved as: | |
| `output/image/image_seg.nii`, if `False`, saving as `output/image_seg.nii`. Default to `True`. | |
| print_log: whether to print logs when saving. Default to `True`. | |
| output_format: an optional string of filename extension to specify the output image writer. | |
| see also: `monai.data.image_writer.SUPPORTED_WRITERS`. | |
| writer: a customised `monai.data.ImageWriter` subclass to save data arrays. | |
| if `None`, use the default writer from `monai.data.image_writer` according to `output_ext`. | |
| if it's a string, it's treated as a class name or dotted path (such as ``"monai.data.ITKWriter"``); | |
| the supported built-in writer classes are ``"NibabelWriter"``, ``"ITKWriter"``, ``"PILWriter"``. | |
| channel_dim: the index of the channel dimension. Default to `0`. | |
| `None` to indicate no channel dimension. | |
| """ | |
| def __init__( | |
| self, | |
| output_dir: PathLike = "./", | |
| output_postfix: str = "trans", | |
| output_ext: str = ".nii.gz", | |
| output_dtype: DtypeLike = np.float32, | |
| resample: bool = True, | |
| mode: str = "nearest", | |
| padding_mode: str = GridSamplePadMode.BORDER, | |
| scale: Optional[int] = None, | |
| dtype: DtypeLike = np.float64, | |
| squeeze_end_dims: bool = True, | |
| data_root_dir: PathLike = "", | |
| separate_folder: bool = True, | |
| print_log: bool = True, | |
| output_format: str = "", | |
| writer: Union[Type[image_writer.ImageWriter], str, None] = None, | |
| channel_dim: Optional[int] = 0, | |
| ) -> None: | |
| self.folder_layout = FolderLayout( | |
| output_dir=output_dir, | |
| postfix=output_postfix, | |
| extension=output_ext, | |
| parent=separate_folder, | |
| makedirs=True, | |
| data_root_dir=data_root_dir, | |
| ) | |
| self.output_ext = output_ext.lower() or output_format.lower() | |
| if isinstance(writer, str): | |
| writer_, has_built_in = optional_import("monai.data", name=f"{writer}") # search built-in | |
| if not has_built_in: | |
| writer_ = locate(f"{writer}") # search dotted path | |
| if writer_ is None: | |
| raise ValueError(f"writer {writer} not found") | |
| writer = writer_ | |
| self.writers = image_writer.resolve_writer(self.output_ext) if writer is None else (writer,) | |
| self.writer_obj = None | |
| _output_dtype = output_dtype | |
| if self.output_ext == ".png" and _output_dtype not in (np.uint8, np.uint16): | |
| _output_dtype = np.uint8 | |
| if self.output_ext == ".dcm" and _output_dtype not in (np.uint8, np.uint16): | |
| _output_dtype = np.uint8 | |
| self.init_kwargs = {"output_dtype": _output_dtype, "scale": scale} | |
| self.data_kwargs = {"squeeze_end_dims": squeeze_end_dims, "channel_dim": channel_dim} | |
| self.meta_kwargs = {"resample": resample, "mode": mode, "padding_mode": padding_mode, "dtype": dtype} | |
| self.write_kwargs = {"verbose": print_log} | |
| self._data_index = 0 | |
| def set_options(self, init_kwargs=None, data_kwargs=None, meta_kwargs=None, write_kwargs=None): | |
| """ | |
| Set the options for the underlying writer by updating the `self.*_kwargs` dictionaries. | |
| The arguments correspond to the following usage: | |
| - `writer = ImageWriter(**init_kwargs)` | |
| - `writer.set_data_array(array, **data_kwargs)` | |
| - `writer.set_metadata(meta_data, **meta_kwargs)` | |
| - `writer.write(filename, **write_kwargs)` | |
| """ | |
| if init_kwargs is not None: | |
| self.init_kwargs.update(init_kwargs) | |
| if data_kwargs is not None: | |
| self.data_kwargs.update(data_kwargs) | |
| if meta_kwargs is not None: | |
| self.meta_kwargs.update(meta_kwargs) | |
| if write_kwargs is not None: | |
| self.write_kwargs.update(write_kwargs) | |
| def __call__(self, img: Union[torch.Tensor, np.ndarray], meta_data: Optional[Dict] = None): | |
| """ | |
| Args: | |
| img: target data content that save into file. The image should be channel-first, shape: `[C,H,W,[D]]`. | |
| meta_data: key-value pairs of metadata corresponding to the data. | |
| """ | |
| meta_data = img.meta if isinstance(img, MetaTensor) else meta_data | |
| subject = "RTStruct"#meta_data["patient_name"] if meta_data else str(self._data_index) | |
| patch_index = None#meta_data.get(Key.PATCH_INDEX, None) if meta_data else None | |
| filename = self.folder_layout.filename(subject=f"{subject}", idx=patch_index) | |
| if meta_data and len(ensure_tuple(meta_data.get("spatial_shape", ()))) == len(img.shape): | |
| self.data_kwargs["channel_dim"] = None | |
| err = [] | |
| for writer_cls in self.writers: | |
| try: | |
| writer_obj = writer_cls(**self.init_kwargs) | |
| writer_obj.set_data_array(data_array=img, **self.data_kwargs) | |
| writer_obj.set_metadata(meta_dict=meta_data, **self.meta_kwargs) | |
| writer_obj.write(filename, **self.write_kwargs) | |
| self.writer_obj = writer_obj | |
| except Exception as e: | |
| print('err',e) | |
| else: | |
| self._data_index += 1 | |
| return img | |
| msg = "\n".join([f"{e}" for e in err]) | |
| raise RuntimeError( | |
| f"{self.__class__.__name__} cannot find a suitable writer for {filename}.\n" | |
| " Please install the writer libraries, see also the installation instructions:\n" | |
| " https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies.\n" | |
| f" The current registered writers for {self.output_ext}: {self.writers}.\n{msg}" | |
| ) | |