-
Notifications
You must be signed in to change notification settings - Fork 56
/
rtstruct.py
142 lines (116 loc) · 4.64 KB
/
rtstruct.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
from typing import List, Union
import numpy as np
from pydicom.dataset import FileDataset
from rt_utils.utils import ROIData
from . import ds_helper, image_helper
class RTStruct:
"""
Wrapper class to facilitate appending and extracting ROI's within an RTStruct
"""
def __init__(self, series_data, ds: FileDataset, ROIGenerationAlgorithm=0):
self.series_data = series_data
self.ds = ds
self.frame_of_reference_uid = ds.ReferencedFrameOfReferenceSequence[
-1
].FrameOfReferenceUID # Use last strucitured set ROI
def set_series_description(self, description: str):
"""
Set the series description for the RTStruct dataset
"""
self.ds.SeriesDescription = description
def add_roi(
self,
mask: np.ndarray,
color: Union[str, List[int]] = None,
name: str = None,
description: str = "",
use_pin_hole: bool = False,
approximate_contours: bool = True,
roi_generation_algorithm: Union[str, int] = 0,
):
"""
Add a ROI to the rtstruct given a 3D binary mask for the ROI's at each slice
Optionally input a color or name for the ROI
If use_pin_hole is set to true, will cut a pinhole through ROI's with holes in them so that they are represented with one contour
If approximate_contours is set to False, no approximation will be done when generating contour data, leading to much larger amount of contour data
"""
# TODO test if name already exists
self.validate_mask(mask)
roi_number = len(self.ds.StructureSetROISequence) + 1
roi_data = ROIData(
mask,
color,
roi_number,
name,
self.frame_of_reference_uid,
description,
use_pin_hole,
approximate_contours,
roi_generation_algorithm,
)
self.ds.ROIContourSequence.append(
ds_helper.create_roi_contour(roi_data, self.series_data)
)
self.ds.StructureSetROISequence.append(
ds_helper.create_structure_set_roi(roi_data)
)
self.ds.RTROIObservationsSequence.append(
ds_helper.create_rtroi_observation(roi_data)
)
def validate_mask(self, mask: np.ndarray) -> bool:
if mask.dtype != bool:
raise RTStruct.ROIException(
f"Mask data type must be boolean. Got {mask.dtype}"
)
if mask.ndim != 3:
raise RTStruct.ROIException(f"Mask must be 3 dimensional. Got {mask.ndim}")
if len(self.series_data) != np.shape(mask)[2]:
raise RTStruct.ROIException(
"Mask must have the save number of layers (In the 3rd dimension) as input series. "
+ f"Expected {len(self.series_data)}, got {np.shape(mask)[2]}"
)
if np.sum(mask) == 0:
raise RTStruct.ROIException("Mask cannot be empty")
return True
def get_roi_names(self) -> List[str]:
"""
Returns a list of the names of all ROI within the RTStruct
"""
if not self.ds.StructureSetROISequence:
return []
return [
structure_roi.ROIName for structure_roi in self.ds.StructureSetROISequence
]
def get_roi_mask_by_name(self, name) -> np.ndarray:
"""
Returns the 3D binary mask of the ROI with the given input name
"""
for structure_roi in self.ds.StructureSetROISequence:
if structure_roi.ROIName.casefold() == name.casefold():
contour_sequence = ds_helper.get_contour_sequence_by_roi_number(
self.ds, structure_roi.ROINumber
)
return image_helper.create_series_mask_from_contour_sequence(
self.series_data, contour_sequence
)
raise RTStruct.ROIException(f"ROI of name `{name}` does not exist in RTStruct")
def save(self, file_path: str):
"""
Saves the RTStruct with the specified name / location
Automatically adds '.dcm' as a suffix
"""
# Add .dcm if needed
file_path = file_path if file_path.endswith(".dcm") else file_path + ".dcm"
try:
file = open(file_path, "w")
# Opening worked, we should have a valid file_path
print("Writing file to", file_path)
self.ds.save_as(file_path)
file.close()
except OSError:
raise Exception(f"Cannot write to file path '{file_path}'")
class ROIException(Exception):
"""
Exception class for invalid ROI masks
"""
pass