forked from qurit/rt-utils
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_rtstruct_builder.py
253 lines (193 loc) · 8.46 KB
/
test_rtstruct_builder.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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
from rt_utils.rtstruct import RTStruct
import pytest
import os
from rt_utils import RTStructBuilder
from rt_utils.utils import SOPClassUID
from rt_utils import image_helper
from pydicom.dataset import validate_file_meta
import numpy as np
def test_create_from_empty_series_dir():
empty_dir_path = os.path.join(os.path.dirname(__file__), "empty")
assert os.path.exists(empty_dir_path)
with pytest.raises(Exception):
RTStructBuilder.create_new(empty_dir_path)
def test_only_images_loaded_into_series_data(new_rtstruct: RTStruct):
assert len(new_rtstruct.series_data) > 0
for ds in new_rtstruct.series_data:
assert hasattr(ds, "pixel_array")
def test_valid_filemeta(new_rtstruct: RTStruct):
try:
validate_file_meta(new_rtstruct.ds.file_meta)
except Exception:
pytest.fail("Invalid file meta in RTStruct dataset")
def test_add_non_binary_roi(new_rtstruct: RTStruct):
mask = get_empty_mask(new_rtstruct)
mask = mask.astype(float)
with pytest.raises(RTStruct.ROIException):
new_rtstruct.add_roi(mask)
def test_add_empty_roi(new_rtstruct: RTStruct):
mask = get_empty_mask(new_rtstruct)
new_rtstruct.add_roi(mask)
assert len(new_rtstruct.ds.ROIContourSequence) == 1
assert len(new_rtstruct.ds.ROIContourSequence[0].ContourSequence) == 0 # No slices added
assert len(new_rtstruct.ds.StructureSetROISequence) == 1
assert len(new_rtstruct.ds.RTROIObservationsSequence) == 1
def test_add_invalid_sized_roi(new_rtstruct: RTStruct):
mask = get_empty_mask(new_rtstruct)
mask = mask[:, :, 1:] # One less slice than expected
with pytest.raises(RTStruct.ROIException):
new_rtstruct.add_roi(mask)
def test_add_valid_roi(new_rtstruct: RTStruct):
assert new_rtstruct.get_roi_names() == []
assert len(new_rtstruct.ds.ROIContourSequence) == 0
assert len(new_rtstruct.ds.StructureSetROISequence) == 0
assert len(new_rtstruct.ds.RTROIObservationsSequence) == 0
NAME = "Test ROI"
COLOR = [123, 123, 232]
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
new_rtstruct.add_roi(mask, color=COLOR, name=NAME)
assert len(new_rtstruct.ds.ROIContourSequence) == 1
assert (
len(new_rtstruct.ds.ROIContourSequence[0].ContourSequence) == 1
) # Only 1 slice was added
assert len(new_rtstruct.ds.StructureSetROISequence) == 1
assert len(new_rtstruct.ds.RTROIObservationsSequence) == 1
assert new_rtstruct.ds.ROIContourSequence[0].ROIDisplayColor == COLOR
assert new_rtstruct.get_roi_names() == [NAME]
def test_get_invalid_roi_mask_by_name(new_rtstruct: RTStruct):
assert new_rtstruct.get_roi_names() == []
with pytest.raises(RTStruct.ROIException):
new_rtstruct.get_roi_mask_by_name("FAKE_NAME")
def test_loading_invalid_rt_struct(series_path):
invalid_rt_struct_path = os.path.join(series_path, "ct_1.dcm")
assert os.path.exists(invalid_rt_struct_path)
with pytest.raises(Exception):
RTStructBuilder.create_from(series_path, invalid_rt_struct_path)
def test_loading_invalid_reference_rt_struct(series_path):
# This RTStruct references images not found within the series path
invalid_reference_rt_struct_path = os.path.join(
series_path, "invalid_reference_rt.dcm"
)
assert os.path.exists(invalid_reference_rt_struct_path)
with pytest.raises(Exception):
RTStructBuilder.create_from(series_path, invalid_reference_rt_struct_path)
def test_non_existant_referenced_study_sequence(series_path):
non_existent_reference_study_rt_struct_path = os.path.join(
series_path, "non_existent_reference_rt.dcm"
)
assert os.path.exists(non_existent_reference_study_rt_struct_path)
rtstruct = RTStructBuilder.create_from(
series_path, non_existent_reference_study_rt_struct_path
)
# Test that the attribute does not exist but RTStruct instantiation was still successful
assert not hasattr(
rtstruct.ds.ReferencedFrameOfReferenceSequence[0], "RTReferencedStudySequence"
)
def test_loading_valid_rt_struct(series_path):
valid_rt_struct_path = os.path.join(series_path, "rt.dcm")
assert os.path.exists(valid_rt_struct_path)
rtstruct = RTStructBuilder.create_from(series_path, valid_rt_struct_path)
# Tests existing values predefined in the file are found
assert hasattr(rtstruct.ds, "ROIContourSequence")
assert hasattr(rtstruct.ds, "StructureSetROISequence")
assert hasattr(rtstruct.ds, "RTROIObservationsSequence")
assert len(rtstruct.ds.ROIContourSequence) == 1
assert len(rtstruct.ds.StructureSetROISequence) == 1
assert len(rtstruct.ds.RTROIObservationsSequence) == 1
# Test adding a new ROI
mask = get_empty_mask(rtstruct)
mask[50:100, 50:100, 0] = 1
rtstruct.add_roi(mask)
assert len(rtstruct.ds.ROIContourSequence) == 2 # 1 should be added
assert len(rtstruct.ds.StructureSetROISequence) == 2 # 1 should be added
assert len(rtstruct.ds.RTROIObservationsSequence) == 2 # 1 should be added
new_roi = rtstruct.ds.StructureSetROISequence[-1]
assert new_roi.ROIName == "ROI-2"
def test_loaded_mask_iou(new_rtstruct: RTStruct):
# Put weird shape in mask
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
mask[60:150, 40:120, 0] = 1
IOU_threshold = 1.0 # Expected accuracy
run_mask_iou_test(new_rtstruct, mask, IOU_threshold)
def test_mask_with_holes_iou(new_rtstruct: RTStruct):
# Create square mask with holes
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
mask[65:85, 65:85, 0] = 0
# add another structure with hole inside of hole
mask[120:230, 120:230, 0] = 1
mask[135:215, 135:215, 0] = 0
mask[150:200, 150:200, 0] = 1
mask[165:185, 165:185, 0] = 0
IOU_threshold = 1
run_mask_iou_test(new_rtstruct, mask, IOU_threshold)
def test_pin_hole_iou(new_rtstruct: RTStruct):
# Create square mask with hole
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
mask[65:85, 65:85, 0] = 0
IOU_threshold = 0.95 # Expect lower accuracy holes lose information
run_mask_iou_test(new_rtstruct, mask, IOU_threshold, use_pin_hole=True)
def test_no_approximation_iou(new_rtstruct: RTStruct):
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
mask[60:150, 40:120, 0] = 1
IOU_threshold = 1.0 # Expected accuracy
run_mask_iou_test(new_rtstruct, mask, IOU_threshold, approximate_contours=False)
def test_contour_data_sizes(new_rtstruct: RTStruct):
mask = get_empty_mask(new_rtstruct)
mask[50:100, 50:100, 0] = 1
mask[60:150, 40:120, 0] = 1
# Given we've added the same mask with and without contour approximation
new_rtstruct.add_roi(mask)
new_rtstruct.add_roi(mask, approximate_contours=False)
# Then using approximation leads to less data within the contour data
assert get_data_len_by_index(new_rtstruct, 0) < get_data_len_by_index(
new_rtstruct, 1
)
def test_nonstandard_image_orientation(oriented_rtstruct: RTStruct):
mask = get_empty_mask(oriented_rtstruct)
mask[10:70, 5:15, 1] = 1
mask[60:70, 5:40, 1] = 1
IOU_threshold = 1.0 # Expected accuracy
run_mask_iou_test(oriented_rtstruct, mask, IOU_threshold)
def test_one_slice_image(one_slice_rtstruct: RTStruct):
mask = get_empty_mask(one_slice_rtstruct)
mask[10:70, 5:15, 0] = 1
mask[60:70, 5:40, 0] = 1
IOU_threshold = 1.0 # Expected accuracy
run_mask_iou_test(one_slice_rtstruct, mask, IOU_threshold)
def get_data_len_by_index(rt_struct: RTStruct, i: int):
return len(rt_struct.ds.ROIContourSequence[i].ContourSequence[0].ContourData)
def run_mask_iou_test(
rtstruct: RTStruct,
mask,
IOU_threshold,
use_pin_hole=False,
approximate_contours=True,
):
# Save and load mask
mask_name = "test"
rtstruct.add_roi(
mask,
name=mask_name,
use_pin_hole=use_pin_hole,
approximate_contours=approximate_contours,
)
loaded_mask = rtstruct.get_roi_mask_by_name(mask_name)
# Use IOU to test accuracy of loaded mask
numerator = np.logical_and(mask, loaded_mask)
denominator = np.logical_or(mask, loaded_mask)
IOU = np.sum(numerator) / np.sum(denominator)
assert IOU >= IOU_threshold
def get_empty_mask(rtstruct) -> np.ndarray:
ref_dicom_image = rtstruct.series_data[0]
mask_dims = (
int(ref_dicom_image.Columns),
int(ref_dicom_image.Rows),
len(rtstruct.series_data),
)
mask = np.zeros(mask_dims)
return mask.astype(bool)