Spaces:
Sleeping
Sleeping
Changed thresholding/resampling order step
Browse files- lungtumormask/dataprocessing.py +11 -16
lungtumormask/dataprocessing.py
CHANGED
@@ -8,6 +8,7 @@ import torch
|
|
8 |
import numpy as np
|
9 |
from monai.transforms import (Compose, LoadImaged, ToNumpyd, ThresholdIntensityd, AddChanneld, NormalizeIntensityd, SpatialCropd, DivisiblePadd, Spacingd, SqueezeDimd)
|
10 |
from tqdm import tqdm
|
|
|
11 |
|
12 |
def mask_lung(scan_path, batch_size=20):
|
13 |
model = lungmask.mask.get_model('unet', 'R231')
|
@@ -50,7 +51,6 @@ def mask_lung(scan_path, batch_size=20):
|
|
50 |
|
51 |
outmask = lungmask.utils.postprocessing(timage_res)
|
52 |
|
53 |
-
|
54 |
outmask = np.asarray(
|
55 |
[lungmask.utils.reshape_mask(outmask[i], xnew_box[i], inimg_raw.shape[1:]) for i in range(outmask.shape[0])],
|
56 |
dtype=np.uint8)
|
@@ -61,7 +61,6 @@ def mask_lung(scan_path, batch_size=20):
|
|
61 |
return outmask.astype(np.uint8), scan_read['image_meta_dict']['affine']
|
62 |
|
63 |
def calculate_extremes(image, annotation_value):
|
64 |
-
|
65 |
holder = np.copy(image)
|
66 |
|
67 |
x_min = float('inf')
|
@@ -98,7 +97,6 @@ def calculate_extremes(image, annotation_value):
|
|
98 |
return ((x_min, x_max), (y_min, y_max), (z_min, z_max))
|
99 |
|
100 |
def process_lung_scan(scan_dict, extremes):
|
101 |
-
|
102 |
load_transformer = Compose(
|
103 |
[
|
104 |
LoadImaged(keys=["image"]),
|
@@ -123,15 +121,12 @@ def process_lung_scan(scan_dict, extremes):
|
|
123 |
)
|
124 |
|
125 |
processed_2 = transformer_1(processed_1)
|
126 |
-
|
127 |
affine = processed_2['image_meta_dict']['affine']
|
128 |
-
|
129 |
normalized_image = processed_2['image']
|
130 |
|
131 |
return normalized_image, affine
|
132 |
|
133 |
def preprocess(image_path):
|
134 |
-
|
135 |
preprocess_dump = {}
|
136 |
|
137 |
scan_dict = {
|
@@ -208,7 +203,6 @@ def find_pad_edge(original):
|
|
208 |
|
209 |
return a_min, a_max + 1, b_min, b_max + 1, c_min, c_max + 1
|
210 |
|
211 |
-
|
212 |
def remove_pad(mask, original):
|
213 |
a_min, a_max, b_min, b_max, c_min, c_max = find_pad_edge(original)
|
214 |
|
@@ -216,27 +210,25 @@ def remove_pad(mask, original):
|
|
216 |
|
217 |
def voxel_space(image, target):
|
218 |
image = Resize((target[0][1]-target[0][0], target[1][1]-target[1][0], target[2][1]-target[2][0]), mode='trilinear')(np.expand_dims(image, 0))[0]
|
219 |
-
image = ThresholdIntensity(above = False, threshold = 0.5, cval = 1)(image)
|
220 |
-
image = ThresholdIntensity(above = True, threshold = 0.5, cval = 0)(image)
|
221 |
|
222 |
return image
|
223 |
|
224 |
def stitch(org_shape, cropped, roi):
|
225 |
-
holder = np.zeros(org_shape)
|
226 |
holder[roi[0][0]:roi[0][1], roi[1][0]:roi[1][1], roi[2][0]:roi[2][1]] = cropped
|
227 |
|
228 |
return holder
|
229 |
|
230 |
-
def post_process(
|
231 |
-
|
232 |
-
|
233 |
-
|
234 |
-
left = remove_pad(left_mask, preprocess_dump['left_lung'].squeeze(0).squeeze(0).numpy())
|
235 |
-
right = remove_pad(right_mask, preprocess_dump['right_lung'].squeeze(0).squeeze(0).numpy())
|
236 |
|
237 |
left = voxel_space(left, preprocess_dump['left_extremes'])
|
238 |
right = voxel_space(right, preprocess_dump['right_extremes'])
|
239 |
|
|
|
|
|
|
|
240 |
left = stitch(preprocess_dump['org_shape'], left, preprocess_dump['left_extremes'])
|
241 |
right = stitch(preprocess_dump['org_shape'], right, preprocess_dump['right_extremes'])
|
242 |
|
@@ -245,5 +237,8 @@ def post_process(left_mask, right_mask, preprocess_dump, lung_filter, threshold)
|
|
245 |
# filter tumor predictions outside the predicted lung area
|
246 |
if lung_filter:
|
247 |
stitched[preprocess_dump['lungmask'] == 0] = 0
|
|
|
|
|
|
|
248 |
|
249 |
return stitched
|
|
|
8 |
import numpy as np
|
9 |
from monai.transforms import (Compose, LoadImaged, ToNumpyd, ThresholdIntensityd, AddChanneld, NormalizeIntensityd, SpatialCropd, DivisiblePadd, Spacingd, SqueezeDimd)
|
10 |
from tqdm import tqdm
|
11 |
+
from skimage.morphology import binary_closing, ball
|
12 |
|
13 |
def mask_lung(scan_path, batch_size=20):
|
14 |
model = lungmask.mask.get_model('unet', 'R231')
|
|
|
51 |
|
52 |
outmask = lungmask.utils.postprocessing(timage_res)
|
53 |
|
|
|
54 |
outmask = np.asarray(
|
55 |
[lungmask.utils.reshape_mask(outmask[i], xnew_box[i], inimg_raw.shape[1:]) for i in range(outmask.shape[0])],
|
56 |
dtype=np.uint8)
|
|
|
61 |
return outmask.astype(np.uint8), scan_read['image_meta_dict']['affine']
|
62 |
|
63 |
def calculate_extremes(image, annotation_value):
|
|
|
64 |
holder = np.copy(image)
|
65 |
|
66 |
x_min = float('inf')
|
|
|
97 |
return ((x_min, x_max), (y_min, y_max), (z_min, z_max))
|
98 |
|
99 |
def process_lung_scan(scan_dict, extremes):
|
|
|
100 |
load_transformer = Compose(
|
101 |
[
|
102 |
LoadImaged(keys=["image"]),
|
|
|
121 |
)
|
122 |
|
123 |
processed_2 = transformer_1(processed_1)
|
|
|
124 |
affine = processed_2['image_meta_dict']['affine']
|
|
|
125 |
normalized_image = processed_2['image']
|
126 |
|
127 |
return normalized_image, affine
|
128 |
|
129 |
def preprocess(image_path):
|
|
|
130 |
preprocess_dump = {}
|
131 |
|
132 |
scan_dict = {
|
|
|
203 |
|
204 |
return a_min, a_max + 1, b_min, b_max + 1, c_min, c_max + 1
|
205 |
|
|
|
206 |
def remove_pad(mask, original):
|
207 |
a_min, a_max, b_min, b_max, c_min, c_max = find_pad_edge(original)
|
208 |
|
|
|
210 |
|
211 |
def voxel_space(image, target):
|
212 |
image = Resize((target[0][1]-target[0][0], target[1][1]-target[1][0], target[2][1]-target[2][0]), mode='trilinear')(np.expand_dims(image, 0))[0]
|
|
|
|
|
213 |
|
214 |
return image
|
215 |
|
216 |
def stitch(org_shape, cropped, roi):
|
217 |
+
holder = np.zeros(org_shape, dtype="float32")
|
218 |
holder[roi[0][0]:roi[0][1], roi[1][0]:roi[1][1], roi[2][0]:roi[2][1]] = cropped
|
219 |
|
220 |
return holder
|
221 |
|
222 |
+
def post_process(left, right, preprocess_dump, lung_filter, threshold):
|
223 |
+
left = remove_pad(left, preprocess_dump['left_lung'].squeeze(0).squeeze(0).numpy())
|
224 |
+
right = remove_pad(right, preprocess_dump['right_lung'].squeeze(0).squeeze(0).numpy())
|
|
|
|
|
|
|
225 |
|
226 |
left = voxel_space(left, preprocess_dump['left_extremes'])
|
227 |
right = voxel_space(right, preprocess_dump['right_extremes'])
|
228 |
|
229 |
+
left = (left >= threshold).astype(int)
|
230 |
+
right = (right >= threshold).astype(int)
|
231 |
+
|
232 |
left = stitch(preprocess_dump['org_shape'], left, preprocess_dump['left_extremes'])
|
233 |
right = stitch(preprocess_dump['org_shape'], right, preprocess_dump['right_extremes'])
|
234 |
|
|
|
237 |
# filter tumor predictions outside the predicted lung area
|
238 |
if lung_filter:
|
239 |
stitched[preprocess_dump['lungmask'] == 0] = 0
|
240 |
+
|
241 |
+
# final post-processing - fix fragmentation
|
242 |
+
stitched = binary_closing(stitched, footprint=ball(radius=5))
|
243 |
|
244 |
return stitched
|