-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsegment.py
143 lines (133 loc) · 6.18 KB
/
segment.py
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
from pathlib import Path
from typing import Union
from ccitk.common.resource import CineImages, PhaseImage
from ccitk.cmr_segment.common.constants import DEFAULT_MODEL_PATH
from ccitk.cmr_segment.segmentor.torch import TorchSegmentor
from ccitk.cmr_segment.segmentor import CineSegmentor
def segment_4d_image(
image_path: Union[Path, str], output_dir: Union[Path, str], model_path: Union[Path, str] = None,
auto_contrast: bool = False, resample: bool = False, enlarge: bool = False,
overwrite: bool = False, device: Union[int, str] = "cpu"
) -> Path:
"""
Segment 4D cardiac nifti image using 3D UNet
:param image_path: string, or Path
The path of a 4D cine image, or a directory containing nifti of all the phases
:param output_dir: string, or Path
:param model_path: string or Path, optional
The path of the trained model. If not provided, default model will be used, which was trained on UKBB data.
:param auto_contrast: bool
:param resample: bool
:param enlarge: bool
:param overwrite: boolean
If overwrite is False, then the segmentor will not be run if the output path exists.
If overwrite is set to True, then the segmentor will overwrite to the output path, if exists.
:param device: int, or str, optional
By default, we will use CPU. If an integer n is provided, then gpu:n will be used.
:return:
Path, the output location.
"""
import mirtk
if isinstance(image_path, str):
image_path = Path(image_path)
if not image_path.exists():
raise FileNotFoundError(f"{str(image_path)} does not exist.")
if not image_path.is_dir():
if auto_contrast:
contrasted_nii_path = image_path.parent.joinpath("contrasted_4d_image.nii.gz")
mirtk.auto_contrast(str(image_path), str(contrasted_nii_path))
image_path = contrasted_nii_path
mirtk.split_volume(
str(image_path),
str(output_dir.joinpath("sequence", "image_")),
"-sequence",
)
cine = CineImages.from_dir(output_dir.joinpath("sequence"))
else:
cine = CineImages.from_dir(image_path)
if auto_contrast:
output_dir.parent.joinpath("contrasted").mkdir(exist_ok=True, parents=True)
for image in cine:
contrasted_nii_path = output_dir.parent.joinpath("contrasted").joinpath(image.name)
mirtk.auto_contrast(str(image_path), str(contrasted_nii_path))
cine = CineImages.from_dir(output_dir.parent.joinpath("contrasted"))
images = []
for image in cine:
image_path = image.path
if resample:
output_dir.parent.joinpath("resampled").mkdir(parents=True, exist_ok=True)
mirtk.resample_image(
str(image.path),
str(output_dir.parent.joinpath("contrasted").joinpath(image.name)),
'-size', 1.25, 1.25, 2
)
image_path = output_dir.parent.joinpath("contrasted").joinpath(image.name)
if enlarge:
output_dir.parent.joinpath("enlarged").mkdir(parents=True, exist_ok=True)
mirtk.enlarge_image(
str(image.path),
str(output_dir.parent.joinpath("enlarged").joinpath(image.name)),
z=20, value=0
)
image_path = output_dir.parent.joinpath("enlarged").joinpath(image.name)
images.append(PhaseImage(path=image_path, phase=image.phase))
cine = CineImages(images)
if isinstance(output_dir, str):
output_dir = Path(output_dir)
if not output_dir.is_dir():
raise ValueError(f"{str(output_dir)} is not a directory.")
if isinstance(model_path, str):
model_path = Path(model_path)
if model_path is None:
model_path = DEFAULT_MODEL_PATH
else:
if not model_path.exists():
raise FileNotFoundError(f"Model file {str(model_path)} does not exist")
segmentor = TorchSegmentor(model_path=model_path, overwrite=overwrite, device=device)
cine_segmentor = CineSegmentor(phase_segmentor=segmentor)
cine_segmentor.apply(
cine=cine, output_dir=output_dir, overwrite=overwrite
)
return output_dir
def segment_image(
image_path: Union[Path, str], output_path: Union[Path, str] = None, model_path: Union[Path, str] = None,
overwrite: bool = False, device: Union[int, str] = "cpu"
) -> Path:
"""
Segment 3D cardiac nifti image using 3D UNet
:param image_path: string, or Path
The path of the nifti image
:param output_path: string, or Path, optional
If not provided, the segmentation will be saved to the directory of image_path, under filename "seg.nii.gz"
:param model_path: string or Path, optional
The path of the trained model. If not provided, default model will be used, which was trained on UKBB data.
:param overwrite: boolean
If overwrite is False, then the segmentor will not be run if the output path exists.
If overwrite is set to True, then the segmentor will overwrite to the output path, if exists.
:param device: int, or str, optional
By default, we will use CPU. If an integer n is provided, then gpu:n will be used.
:return:
Path, the output location.
"""
segmentor = TorchSegmentor(model_path=model_path, overwrite=overwrite, device=device)
if isinstance(image_path, str):
image_path = Path(image_path)
if not image_path.exists():
raise FileNotFoundError(f"Image file {str(image_path)} does not exist")
if isinstance(output_path, str):
output_path = Path(output_path)
if isinstance(model_path, str):
model_path = Path(model_path)
if model_path is None:
model_path = DEFAULT_MODEL_PATH
else:
if not model_path.exists():
raise FileNotFoundError(f"Model file {str(model_path)} does not exist")
if output_path is None:
output_path = image_path.parent.joinpath("seg.nii.gz")
if not output_path.exists() or overwrite:
segmentor = TorchSegmentor(model_path=model_path, overwrite=overwrite,device=device)
segmentor.execute(
phase_path=image_path, output_path=output_path
)
return output_path