mrid-python 0.1.3__tar.gz → 0.1.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (44) hide show
  1. mrid_python-0.1.4/PKG-INFO +77 -0
  2. mrid_python-0.1.4/README.md +63 -0
  3. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/atlas/MNI152/__init__.py +7 -0
  4. mrid_python-0.1.4/mrid/preprocessing/CTseg.py +81 -0
  5. mrid_python-0.1.4/mrid/preprocessing/__init__.py +12 -0
  6. mrid_python-0.1.4/mrid/preprocessing/hd_bet.py +242 -0
  7. mrid_python-0.1.4/mrid/preprocessing/mask.py +42 -0
  8. mrid_python-0.1.3/mrid/preprocessing/registration.py → mrid_python-0.1.4/mrid/preprocessing/simple_elastix.py +20 -82
  9. mrid_python-0.1.4/mrid/preprocessing/spatial.py +81 -0
  10. mrid_python-0.1.4/mrid/preprocessing/synthstrip.py +259 -0
  11. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/study.py +223 -36
  12. mrid_python-0.1.4/mrid/training/__init__.py +0 -0
  13. mrid_python-0.1.4/mrid/training/slicer.py +235 -0
  14. mrid_python-0.1.4/mrid/training/transforms.py +98 -0
  15. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/__init__.py +1 -0
  16. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/plotting.py +5 -3
  17. mrid_python-0.1.4/mrid_python.egg-info/PKG-INFO +77 -0
  18. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid_python.egg-info/SOURCES.txt +9 -2
  19. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid_python.egg-info/top_level.txt +1 -0
  20. {mrid_python-0.1.3 → mrid_python-0.1.4}/pyproject.toml +1 -1
  21. mrid_python-0.1.3/PKG-INFO +0 -140
  22. mrid_python-0.1.3/README.md +0 -126
  23. mrid_python-0.1.3/mrid/preprocessing/__init__.py +0 -12
  24. mrid_python-0.1.3/mrid/preprocessing/skullstripping.py +0 -185
  25. mrid_python-0.1.3/mrid_python.egg-info/PKG-INFO +0 -140
  26. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/__init__.py +0 -0
  27. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/atlas/SRI24/__init__.py +0 -0
  28. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/atlas/__init__.py +0 -0
  29. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/loading/__init__.py +0 -0
  30. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/loading/convert.py +0 -0
  31. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/preprocessing/bias_field_correction.py +0 -0
  32. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/preprocessing/cropping.py +0 -0
  33. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/dcm2niix.py +0 -0
  34. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/dicom_uid_fixer.py +0 -0
  35. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/python_utils.py +0 -0
  36. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/stl_utils.py +0 -0
  37. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid/utils/torch_utils.py +0 -0
  38. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid_python.egg-info/dependency_links.txt +0 -0
  39. {mrid_python-0.1.3 → mrid_python-0.1.4}/mrid_python.egg-info/requires.txt +0 -0
  40. {mrid_python-0.1.3 → mrid_python-0.1.4}/setup.cfg +0 -0
  41. {mrid_python-0.1.3 → mrid_python-0.1.4}/tests/test_loading.py +0 -0
  42. {mrid_python-0.1.3 → mrid_python-0.1.4}/tests/test_preprocessing.py +0 -0
  43. {mrid_python-0.1.3 → mrid_python-0.1.4}/tests/test_study.py +0 -0
  44. {mrid_python-0.1.3 → mrid_python-0.1.4}/tests/test_utils.py +0 -0
@@ -0,0 +1,77 @@
1
+ Metadata-Version: 2.4
2
+ Name: mrid-python
3
+ Version: 0.1.4
4
+ Summary: Tools for working with 3D medical images and segmentations - registration, brain skull-stripping, etc.
5
+ Author-email: Ivan Nikishev <nkshv2@gmail.com>
6
+ Project-URL: Homepage, https://github.com/inikishev/mrid
7
+ Project-URL: Repository, https://github.com/inikishev/mrid
8
+ Project-URL: Issues, https://github.com/inikishev/mrid/isses
9
+ Keywords: MRI,medical imaging
10
+ Requires-Python: >=3.10
11
+ Description-Content-Type: text/markdown
12
+ Requires-Dist: numpy
13
+ Requires-Dist: SimpleITK
14
+
15
+ <h1 align='center'>mrid</h1>
16
+
17
+ mrid is a library for preprocessing of 3D images, particularly medical images.
18
+
19
+ It provide interfaces for many medical image processing tools such as [SimpleElastix](https://simpleelastix.github.io/), [HD-BET](https://github.com/MIC-DKFZ/HD-BET#Installation), [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/), [CTSeg](https://github.com/WCHN/CTseg). Note that those libraries are not bundled with mrid, I've included installation instructions in all notebooks.
20
+
21
+ ### Installation
22
+
23
+ Either run
24
+
25
+ ```
26
+ pip install mrid-python
27
+ ```
28
+
29
+ or
30
+
31
+ ```
32
+ pip install git+https://github.com/inikishev/mrid
33
+ ```
34
+
35
+ ### Basics
36
+
37
+ The images you pass to all functions in mrid can be path to a .nii.gz file, DICOM directory, sitk.Image, numpy array or torch tensor. All functions return results as `sitk.Image`. If you need a numpy array, you can use `mrid.tonumpy(sitk_image)`.
38
+
39
+ ### Registering images with SimpleITK-SimpleElastix
40
+
41
+ [SimpleElastix](https://simpleelastix.github.io/) is a robust tool for image registration which works really well out-of-the-box. It works on both Windows and Linux.
42
+
43
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/SimpleElastix%20tutorial.ipynb) for how to install and use it.
44
+ <img width="828" height="839" alt="image" src="https://github.com/user-attachments/assets/f083178a-82f0-411d-9d46-ffff774248e0" />
45
+
46
+ ### Skullstripping MRI scans with HD-BET
47
+
48
+ [HD-BET](https://github.com/MIC-DKFZ/HD-BET) is a model that performs skullstripping of pre- and post-constrast T1, T2 and FALIR MRIs. It works on both Windows and Linux.
49
+
50
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/HD-BET%20tutorial.ipynb) for how to install and use it
51
+ <img width="828" height="840" alt="image" src="https://github.com/user-attachments/assets/ec3a8a39-0554-419f-9df9-b8e5bebc9232" />
52
+
53
+ ### Skullstripping with SynthStrip
54
+
55
+ [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/) is a skull-stripping tool that works with many different image types and modalities, including MRI, DWI, CT, PET, etc.
56
+
57
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/SynthStrip%20tutorial.ipynb) for how to install and use it
58
+ <img width="828" height="840" alt="image" src="https://github.com/user-attachments/assets/60fdca28-054d-4e62-93f1-6b84daa3fb9a" />
59
+
60
+ ### Skullstripping and segmentation of CT images with CTseg
61
+
62
+ [CTseg](https://github.com/WCHN/CTseg) can skull-strip CT images and perform their segmentation, it also registers them to a common space (see its README). Note that it can be very slow for 512x512 series (can take few hours), but you can downsample to 256x256. If you only need to quickly skullstrip CT scans without warping them you can use SynthStrip.
63
+
64
+ TODO!!!
65
+
66
+ ### Example workflow - preprocessing MRIs to BraTS format
67
+
68
+ Many [BraTS](https://www.synapse.org/brats) datasets are provided as skullstripped images in SRI24 space. See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/BraTS%20preprocessing%20workflow.ipynb) for how to process raw scans to this format.
69
+
70
+ <img width="828" height="849" alt="image" src="https://github.com/user-attachments/assets/f1b38db3-6648-4660-a381-d68a2eb8508d" />
71
+
72
+ (T1n image looks weird because that's just how it is in the dataset)
73
+
74
+ ### References
75
+ The MRIs for all images above are from https://zenodo.org/records/7213153.
76
+
77
+ > Colin Vanden Bulcke. (2022). Open-Access DICOM MRI session (1.0) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.7213153
@@ -0,0 +1,63 @@
1
+ <h1 align='center'>mrid</h1>
2
+
3
+ mrid is a library for preprocessing of 3D images, particularly medical images.
4
+
5
+ It provide interfaces for many medical image processing tools such as [SimpleElastix](https://simpleelastix.github.io/), [HD-BET](https://github.com/MIC-DKFZ/HD-BET#Installation), [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/), [CTSeg](https://github.com/WCHN/CTseg). Note that those libraries are not bundled with mrid, I've included installation instructions in all notebooks.
6
+
7
+ ### Installation
8
+
9
+ Either run
10
+
11
+ ```
12
+ pip install mrid-python
13
+ ```
14
+
15
+ or
16
+
17
+ ```
18
+ pip install git+https://github.com/inikishev/mrid
19
+ ```
20
+
21
+ ### Basics
22
+
23
+ The images you pass to all functions in mrid can be path to a .nii.gz file, DICOM directory, sitk.Image, numpy array or torch tensor. All functions return results as `sitk.Image`. If you need a numpy array, you can use `mrid.tonumpy(sitk_image)`.
24
+
25
+ ### Registering images with SimpleITK-SimpleElastix
26
+
27
+ [SimpleElastix](https://simpleelastix.github.io/) is a robust tool for image registration which works really well out-of-the-box. It works on both Windows and Linux.
28
+
29
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/SimpleElastix%20tutorial.ipynb) for how to install and use it.
30
+ <img width="828" height="839" alt="image" src="https://github.com/user-attachments/assets/f083178a-82f0-411d-9d46-ffff774248e0" />
31
+
32
+ ### Skullstripping MRI scans with HD-BET
33
+
34
+ [HD-BET](https://github.com/MIC-DKFZ/HD-BET) is a model that performs skullstripping of pre- and post-constrast T1, T2 and FALIR MRIs. It works on both Windows and Linux.
35
+
36
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/HD-BET%20tutorial.ipynb) for how to install and use it
37
+ <img width="828" height="840" alt="image" src="https://github.com/user-attachments/assets/ec3a8a39-0554-419f-9df9-b8e5bebc9232" />
38
+
39
+ ### Skullstripping with SynthStrip
40
+
41
+ [SynthStrip](https://surfer.nmr.mgh.harvard.edu/docs/synthstrip/) is a skull-stripping tool that works with many different image types and modalities, including MRI, DWI, CT, PET, etc.
42
+
43
+ See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/SynthStrip%20tutorial.ipynb) for how to install and use it
44
+ <img width="828" height="840" alt="image" src="https://github.com/user-attachments/assets/60fdca28-054d-4e62-93f1-6b84daa3fb9a" />
45
+
46
+ ### Skullstripping and segmentation of CT images with CTseg
47
+
48
+ [CTseg](https://github.com/WCHN/CTseg) can skull-strip CT images and perform their segmentation, it also registers them to a common space (see its README). Note that it can be very slow for 512x512 series (can take few hours), but you can downsample to 256x256. If you only need to quickly skullstrip CT scans without warping them you can use SynthStrip.
49
+
50
+ TODO!!!
51
+
52
+ ### Example workflow - preprocessing MRIs to BraTS format
53
+
54
+ Many [BraTS](https://www.synapse.org/brats) datasets are provided as skullstripped images in SRI24 space. See [this notebook](https://github.com/inikishev/mrid/blob/main/notebooks/BraTS%20preprocessing%20workflow.ipynb) for how to process raw scans to this format.
55
+
56
+ <img width="828" height="849" alt="image" src="https://github.com/user-attachments/assets/f1b38db3-6648-4660-a381-d68a2eb8508d" />
57
+
58
+ (T1n image looks weird because that's just how it is in the dataset)
59
+
60
+ ### References
61
+ The MRIs for all images above are from https://zenodo.org/records/7213153.
62
+
63
+ > Colin Vanden Bulcke. (2022). Open-Access DICOM MRI session (1.0) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.7213153
@@ -62,6 +62,13 @@ def get_mni152(
62
62
  ):
63
63
  """Returns path to .nii.gz file of specified MNI-152 template.
64
64
 
65
+ The following templates are available:
66
+ - ``"2006 T1w symmetric"``
67
+ - ``"2009a T1w symmetric"``
68
+ - ``"2009a T2w symmetric"``
69
+ - ``"2009a T1w asymmetric"``
70
+ - ``"2009a T2w asymmetric"``
71
+
65
72
  Descriptions of templates are available here https://zenodo.org/records/15470657
66
73
  """
67
74
  filename = f"{type} {bool(skullstripped)}.nii.gz"
@@ -0,0 +1,81 @@
1
+ """"
2
+ This requires CTseg docker image to be present in the system.
3
+
4
+ If you haven't already, install docker to your OS following this https://docs.docker.com/engine/install/
5
+
6
+ CTseg is available in the following repository: https://github.com/WCHN/CTseg
7
+
8
+ Navigate to any folder, open the terminal in that folder and type in ``git clone https://github.com/WCHN/CTseg``. You might have to [install git](https://git-scm.com/install/) if you don't have it installed. This will create a new directory called ``CTseg`` in the folder and download the repository to it. Alternatively you can open the repository in your web browser, click on the green "Code" button near the top and click "Download ZIP", and unpack the archive to a folder named ``CTseg``.
9
+
10
+ Then type in the command specified in "Docker" section of the read-me in https://github.com/WCHN/CTseg to build an image from the Dockerfile in this repository. I decided to not copy the command here in case it gets updated in the repository. This will download and build a docker image, usually to ``/var/lib/docker/`` on Linux, or inside the Docker Desktop virtual machine disk image on Windows. Note that this will download 3 GB of data, so it might take some time.
11
+
12
+ After it is done, the CTseg functions from mrid can be used.
13
+ """
14
+ import os
15
+ import subprocess
16
+ from pathlib import Path
17
+
18
+ from ..loading import tositk
19
+
20
+ def run_CTseg(
21
+ pth_ct: str | os.PathLike,
22
+ dir_out: str = "",
23
+ docker_image = "ubuntu:ctseg",
24
+ ) -> None:
25
+ """Runs ``CTseg`` command-line routine via ``subprocess.run``.
26
+
27
+ Args:
28
+ pth_ct (str | os.PathLike): path to a file which must be in a ``*.nii`` format.
29
+ dir_out (str, optional):
30
+ optional name of a directory that will be created next to ``path_ct`` nii file to save CTseg outputs to.
31
+ If empty, outputs are saved next to ``path_ct`` file. Defaults to ''.
32
+ docker_image (str, optional): name of the docker image that ``CTseg`` is installed in. Defaults to "ubuntu:ctseg".
33
+ """
34
+
35
+ # docker run --rm -it -v "/home/jj/data":/data ubuntu:ctseg function spm_CTseg '/data/CT.nii'
36
+ # better
37
+ # docker run --rm -it -v "/home/jj/data":/data ubuntu:ctseg eval "spm_CTseg('/data/CT.nii', '', true, true, true, true, 1.0)"
38
+ pth_ct = Path(pth_ct)
39
+
40
+ if dir_out != "":
41
+ if "/" in dir_out or "\\" in dir_out:
42
+ raise RuntimeError(
43
+ "dir_out should be name of directory that will be created next to `path_ct`. "
44
+ f"It can't be a path. Got '{dir_out}'"
45
+ )
46
+
47
+ dir_out = f"/data/{dir_out}"
48
+
49
+ command = [
50
+ "docker",
51
+ "run",
52
+
53
+ # --rm automatically removes the container's file system after the container exits. This is useful for running temporary containers.
54
+ "--rm",
55
+
56
+ # This is a combination of two flags, -i and -t:
57
+ # -i (interactive): Keeps the standard input (STDIN) open, allowing you to interact with the container.
58
+ # -t (tty): Allocates a pseudo-TTY, which makes the container behave like a normal terminal session. (doesn't work with subprocess)
59
+ #"-it",
60
+ "-i",
61
+
62
+ # -v used for mounting volumes, which allows you to connect a file path on your host machine to a path inside the container. This option requires a specific format: -v <host_path>:<container_path>.
63
+ "-v",
64
+ f"{os.path.normpath(pth_ct.parent)}:/data",
65
+
66
+ # docker image name
67
+ docker_image,
68
+
69
+ # evaluate matlab code
70
+ "eval",
71
+
72
+ # code to evaluate
73
+ f"spm_CTseg('/data/{pth_ct.name}', '{dir_out}', true, true, true, true, 1.0)",
74
+ ]
75
+
76
+ # run dcm2niix
77
+ subprocess.run(command, check=True)
78
+
79
+ # this creates
80
+ # wc01_1_00001_temp_CT_CTseg.nii
81
+ # wc02_1_00001_temp_CT_CTseg
@@ -0,0 +1,12 @@
1
+ from .bias_field_correction import n4_bias_field_correction
2
+ from .cropping import crop_bg, crop_bg_D
3
+ from .spatial import downsample, resample_to, resize
4
+
5
+ # lib wrappers
6
+ from . import hd_bet, CTseg, simple_elastix, synthstrip, mask
7
+ __all__ = [
8
+ "n4_bias_field_correction",
9
+ "crop_bg", "crop_bg_D",
10
+ "resample_to", "resize", "downsample",
11
+ "hd_bet", "CTseg", "simple_elastix", "synthstrip", "mask",
12
+ ]
@@ -0,0 +1,242 @@
1
+ import os
2
+ import subprocess
3
+ import tempfile
4
+ from collections.abc import Mapping
5
+ from typing import Literal
6
+
7
+ import SimpleITK as sitk
8
+
9
+ from ..loading.convert import ImageLike, tositk
10
+ from ..utils.torch_utils import CUDA_IF_AVAILABLE
11
+ from .simple_elastix import register, register_D
12
+ from .mask import expand_binary_mask, apply_mask
13
+
14
+ # hd_bet -h
15
+
16
+ # -i INPUT, --input INPUT
17
+ # input. Can be either a single file name or an input folder. If file: must be nifti (.nii.gz) and can only be 3D. No support for 4d images, use fslsplit to split 4d sequences into 3d images. If folder: all files ending with .nii.gz within that folder will be brain extracted.
18
+ # -o OUTPUT, --output OUTPUT
19
+ # output. Can be either a filename or a folder. If it does not exist, the folder will be created
20
+ # -device DEVICE used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'. Default: cuda
21
+ # --disable_tta Set this flag to disable test time augmentation. This will make prediction faster at a slight decrease in prediction quality. Recommended for device cpu
22
+ # --save_bet_mask Set this flag to keep the bet masks. Otherwise they will be removed once HD_BET is done
23
+ # --no_bet_image Set this flag to disable generating the skull stripped/brain extracted image. Only makes sense if you also set --save_bet_mask
24
+ # --verbose Talk to me.
25
+
26
+ def run_hd_bet(
27
+ input: str | os.PathLike,
28
+ output: str | os.PathLike,
29
+ device: Literal['cpu', 'cuda', 'mps'] = CUDA_IF_AVAILABLE,
30
+ disable_tta: bool = False,
31
+ save_bet_mask: bool = True,
32
+ no_bet_image: bool = False,
33
+ verbose: bool = False,
34
+ ) -> None:
35
+ """Runs HD-BET command-line routine via ``subprocess.run``.
36
+
37
+ Args:
38
+ input (str | os.PathLike):
39
+ input. Can be either a single file name or an input folder.
40
+ If file: must be nifti (.nii.gz) and can only be 3D.
41
+ No support for 4d images, use fslsplit to split 4d sequences into 3d images.
42
+ If folder: all files ending with .nii.gz within that folder will be brain extracted.
43
+ output (str | os.PathLike):
44
+ output. Can be either a filename or a folder. If it does not exist, the folder will be created
45
+ device (str, optional):
46
+ used to set on which device the prediction will run.
47
+ Can be 'cuda' (=GPU), 'cpu' or 'mps'. Default: cuda. Defaults to CUDA_IF_AVAILABLE.
48
+ disable_tta (bool, optional):
49
+ Set this flag to disable test time augmentation.
50
+ This will make prediction faster at a slight decrease in prediction quality.
51
+ Recommended for device cpu. Defaults to False.
52
+ save_bet_mask (bool, optional):
53
+ Set this flag to keep the bet masks.
54
+ Otherwise they will be removed once HD_BET is done. Defaults to True.
55
+ no_bet_image (bool, optional):
56
+ Set this flag to disable generating the skull stripped/brain extracted image.
57
+ Only makes sense if you also set --save_bet_mask. Defaults to False.
58
+ verbose (bool, optional): Talk to me. Defaults to False.
59
+ """
60
+
61
+ command = [
62
+ "hd-bet",
63
+ "-i", os.path.normpath(input),
64
+ "-o", os.path.normpath(output),
65
+ "-device", device,
66
+ ]
67
+ if disable_tta: command.append("--disable_tta")
68
+ if save_bet_mask: command.append("--save_bet_mask")
69
+ if no_bet_image: command.append("--no_bet_image")
70
+ if verbose: command.append("--verbose")
71
+
72
+ # run dcm2niix
73
+ subprocess.run(command, check=True)
74
+
75
+
76
+ def predict_brain_mask(
77
+ input: ImageLike,
78
+ register_to_mni152: Literal["T1", "T2"] | None = None,
79
+ device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
80
+ disable_tta: bool = False,
81
+ verbose: bool = False,
82
+ ) -> sitk.Image:
83
+ """Returns brain mask of ``input`` predicted by HD-BET.
84
+
85
+ Args:
86
+ input (ImageLike): input to skullstrip. Recommended T1-w, postcontrast T1-w, T2-w or FLAIR sequence in MNI152 space.
87
+ register_to_mni152 (str | None, optional):
88
+ Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
89
+ if specified, ``input`` will be registered to specified MNI152 template,
90
+ then after prediction the brain mask registered back to original ``input``.
91
+ Note that HD-BET expects images to be in MNI152 space. Defaults to None.
92
+ device (str, optional):
93
+ used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
94
+ Defaults to CUDA_IF_AVAILABLE.
95
+ disable_tta (bool, optional):
96
+ Set this flag to disable test time augmentation.
97
+ This will make prediction faster at a slight decrease in prediction quality.
98
+ Recommended for device cpu. Defaults to False.
99
+ verbose (bool, optional): Talk to me. Defaults to False.
100
+ """
101
+ input = tositk(input)
102
+
103
+ # ---------------------------- register to mni152 ---------------------------- #
104
+ if register_to_mni152 is not None:
105
+ from ..atlas.MNI152 import get_mni152
106
+ mni152 = get_mni152(f"2009a {register_to_mni152}w asymmetric", skullstripped=False) # type:ignore
107
+ input_mni = register(input, mni152)
108
+
109
+ else:
110
+ input_mni = input
111
+
112
+ # ---------------------------- predict brain mask ---------------------------- #
113
+ with tempfile.TemporaryDirectory() as tmpdir:
114
+ sitk.WriteImage(input_mni, os.path.join(tmpdir, "input.nii.gz"))
115
+
116
+ run_hd_bet(
117
+ input = os.path.join(tmpdir, "input.nii.gz"),
118
+ output = os.path.join(tmpdir, "output.nii.gz"),
119
+ device=device, disable_tta=disable_tta, save_bet_mask=True, verbose=verbose,
120
+ )
121
+
122
+ brain_mask_mni = tositk(os.path.join(tmpdir, "output_bet.nii.gz"))
123
+
124
+ # ------------------------- unregister mask if needed ------------------------ #
125
+ if register_to_mni152 is not None:
126
+ study_mni = dict(image=input_mni, seg_brain=brain_mask_mni)
127
+ study = register_D(study_mni, key="image", to=input)
128
+ brain_mask = study["seg_brain"]
129
+
130
+ else:
131
+ brain_mask = brain_mask_mni
132
+
133
+ return brain_mask
134
+
135
+ def skullstrip(
136
+ input: ImageLike,
137
+ register_to_mni152: Literal["T1", "T2"] | None = None,
138
+ device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
139
+ disable_tta: bool = False,
140
+ verbose: bool = False,
141
+
142
+ expand: int = 0,
143
+ ) -> sitk.Image:
144
+ """Skullstrips ``input`` using HD-BET.
145
+
146
+ Args:
147
+ input (ImageLike): input to skullstrip. Recommended T1-w, postcontrast T1-w, T2-w or FLAIR sequence in MNI152 space.
148
+ register_to_mni152 (str | None, optional):
149
+ Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
150
+ if specified, ``input`` will be registered to specified MNI152 template,
151
+ then after prediction the brain mask registered back to original ``input``.
152
+ Note that HD-BET expects images to be in MNI152 space. Defaults to None.
153
+ device (str, optional):
154
+ used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
155
+ Defaults to CUDA_IF_AVAILABLE.
156
+ disable_tta (bool, optional):
157
+ Set this flag to disable test time augmentation. This will make prediction faster
158
+ at a slight decrease in prediction quality. Recommended for device cpu. Defaults to False.
159
+ verbose (bool, optional): Talk to me. Defaults to False.
160
+ expand (int, optional):
161
+ Positive values expand brain mask by this many pixels, meaning inner parts of the skull will be included;
162
+ Negative values dilate brain mask by this many pixels, meaning outer parts of the brain will be excluded.
163
+
164
+ """
165
+ input = tositk(input)
166
+ mask = predict_brain_mask(input=input, register_to_mni152=register_to_mni152,
167
+ device=device, disable_tta=disable_tta, verbose=verbose)
168
+
169
+ if expand != 0:
170
+ mask = expand_binary_mask(mask, expand=expand)
171
+
172
+ return apply_mask(input, mask)
173
+
174
+
175
+ def skullstrip_D(
176
+ images: Mapping[str, ImageLike],
177
+ key: str,
178
+ register_to_mni152: Literal["T1", "T2"] | None = None,
179
+ device: Literal["cpu", "cuda", "mps"] = CUDA_IF_AVAILABLE,
180
+ disable_tta: bool = False,
181
+ verbose: bool = False,
182
+
183
+ expand: int = 0,
184
+
185
+ include_mask: bool = False,
186
+ keep_original: bool = False,
187
+ ) -> dict[str, sitk.Image]:
188
+ """Predicts brain mask of ``images[key]`` using HD-BET, then uses this mask to skull strip all values in ``images``.
189
+
190
+ Args:
191
+ images (Mapping[str, ImageLike]): dictionary of images that align with each other.
192
+ key (str): key of the image to pass to HD-BET for brain mask prediction.
193
+ register_to_mni152 (str | None, optional):
194
+ Modality of MNI152 template to pre-register ``input`` to. Should be ``"T1"``, ``"T2"`` or ``None``.
195
+ if specified, ``input`` will be registered to specified MNI152 template,
196
+ then after prediction the brain mask registered back to original ``input``.
197
+ Note that HD-BET expects images to be in MNI152 space. Defaults to None.
198
+ device (str, optional):
199
+ used to set on which device the prediction will run. Can be 'cuda' (=GPU), 'cpu' or 'mps'.
200
+ Defaults to CUDA_IF_AVAILABLE.
201
+ disable_tta (bool, optional):
202
+ Set this flag to disable test time augmentation. This will make prediction faster
203
+ at a slight decrease in prediction quality. Recommended for device cpu. Defaults to False.
204
+ verbose (bool, optional): Talk to me. Defaults to False.
205
+ expand (int, optional):
206
+ Positive values expand brain mask by this many pixels, meaning inner parts of the skull will be included;
207
+ Negative values dilate brain mask by this many pixels, meaning outer parts of the brain will be excluded.
208
+ include_mask (bool, optional):
209
+ if True, adds ``"seg_seg_hd_bet"`` with brain mask predicted by HD-BET to returned dictionary.
210
+ This adds brain mask BEFORE expanding/dilating if ``expand`` argument is specified.
211
+ keep_original (bool, Optional):
212
+ if True, skull-stripped images are added to the dictionary
213
+ with ``"_hd_bet" ``postfix, rather than replacing.
214
+ """
215
+ images = {k: tositk(v) for k,v in images.items()}
216
+
217
+ mask = predict_brain_mask(input=images[key], register_to_mni152=register_to_mni152,
218
+ device=device, disable_tta=disable_tta, verbose=verbose)
219
+
220
+ skullstripped = {}
221
+
222
+ # include mask before expanding
223
+ if include_mask:
224
+ mask_sitk = tositk(mask)
225
+ mask_sitk.CopyInformation(images[key])
226
+ skullstripped["seg_hd_bet"] = mask_sitk
227
+
228
+ # expand
229
+ if expand != 0:
230
+ mask = expand_binary_mask(mask, expand=expand)
231
+
232
+ # apply mask
233
+ for k,v in images.items():
234
+ skullstripped[k] = apply_mask(v, mask)
235
+
236
+ # optionally add with skullstripped postfix
237
+ if keep_original:
238
+ skullstripped = {f"{k}_hd_bet": v for k,v in skullstripped}
239
+ skullstripped.update(images.copy())
240
+
241
+ return skullstripped
242
+
@@ -0,0 +1,42 @@
1
+ import SimpleITK as sitk
2
+ import numpy as np
3
+ import SimpleITK as sitk
4
+
5
+ from ..loading.convert import ImageLike, tositk, tonumpy
6
+
7
+
8
+ def expand_binary_mask(binary_mask: ImageLike, expand: int) -> sitk.Image:
9
+ """Expand or dilate a binary mask.
10
+
11
+ Args:
12
+ binary_mask (ImageLike): mask
13
+ expand (int, optional):
14
+ Positive values expand the mask by this many pixels;
15
+ Negative values dilate the mask by this many pixels.
16
+ """
17
+ binary_mask = tositk(binary_mask)
18
+ if expand > 0:
19
+ inverted_mask = 1 - binary_mask
20
+ return 1 - sitk.BinaryDilate(inverted_mask, (expand, expand, expand))
21
+
22
+ if expand < 0:
23
+ return sitk.BinaryDilate(binary_mask, (-expand, -expand, -expand))
24
+
25
+ return binary_mask
26
+
27
+ def apply_mask(image: ImageLike, mask: ImageLike) -> sitk.Image:
28
+ """Applies ``mask`` to ``image``, that is all values where ``mask > 0`` are kept.
29
+
30
+ This function sets all values outside of the mask to smallest value within the mask."""
31
+ image = tositk(image)
32
+ mask = tositk(mask)
33
+
34
+ mask = sitk.Cast(mask, image.GetPixelID())
35
+
36
+ image_np = sitk.GetArrayFromImage(image)
37
+ mask_np = (tonumpy(mask) > 0).astype(np.bool)
38
+ image_ma = np.ma.masked_array(image_np, ~mask_np)
39
+
40
+ image_applied = tositk(image_ma.filled(image_ma.min()))
41
+ image_applied.CopyInformation(image)
42
+ return image_applied