Spaces:
Sleeping
Sleeping
added more verbose and refactored
Browse files- lungtumormask/dataprocessing.py +5 -4
- lungtumormask/mask.py +1 -1
- lungtumormask/network.py +1 -1
lungtumormask/dataprocessing.py
CHANGED
@@ -7,6 +7,7 @@ from monai.transforms.utility.dictionary import ToTensord
|
|
7 |
import torch
|
8 |
import numpy as np
|
9 |
from monai.transforms import (Compose, LoadImaged, ToNumpyd, ThresholdIntensityd, AddChanneld, NormalizeIntensityd, SpatialCropd, DivisiblePadd, Spacingd, SqueezeDimd)
|
|
|
10 |
|
11 |
def mask_lung(scan_path, batch_size=20):
|
12 |
model = lungmask.mask.get_model('unet', 'R231')
|
@@ -28,7 +29,7 @@ def mask_lung(scan_path, batch_size=20):
|
|
28 |
)
|
29 |
|
30 |
scan_read = transformer(scan_dict)
|
31 |
-
inimg_raw = scan_read['image'].swapaxes(0,2)
|
32 |
|
33 |
tvolslices, xnew_box = lungmask.utils.preprocess(inimg_raw, resolution=[256, 256])
|
34 |
tvolslices[tvolslices > 600] = 600
|
@@ -41,7 +42,7 @@ def mask_lung(scan_path, batch_size=20):
|
|
41 |
timage_res = np.empty((np.append(0, tvolslices[0].shape)), dtype=np.uint8)
|
42 |
|
43 |
with torch.no_grad():
|
44 |
-
for X in dataloader_val:
|
45 |
X = X.float().to(device)
|
46 |
prediction = model(X)
|
47 |
pls = torch.max(prediction, 1)[1].detach().cpu().numpy().astype(np.uint8)
|
@@ -57,7 +58,6 @@ def mask_lung(scan_path, batch_size=20):
|
|
57 |
outmask = np.swapaxes(outmask, 0, 2)
|
58 |
#outmask = np.flip(outmask, 0)
|
59 |
|
60 |
-
|
61 |
return outmask.astype(np.uint8), scan_read['image_meta_dict']['affine']
|
62 |
|
63 |
def calculate_extremes(image, annotation_value):
|
@@ -143,6 +143,7 @@ def preprocess(image_path):
|
|
143 |
preprocess_dump['pixdim'] = im['image_meta_dict']['pixdim'][1:4]
|
144 |
preprocess_dump['org_affine'] = im['image_meta_dict']['affine']
|
145 |
|
|
|
146 |
masked_lungs = mask_lung(image_path, 5)
|
147 |
right_lung_extreme = calculate_extremes(masked_lungs[0], 1)
|
148 |
preprocess_dump['right_extremes'] = right_lung_extreme
|
@@ -249,4 +250,4 @@ def post_process(left_mask, right_mask, preprocess_dump):
|
|
249 |
|
250 |
stitched = np.logical_or(left, right).astype(int)
|
251 |
|
252 |
-
return stitched
|
|
|
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')
|
|
|
29 |
)
|
30 |
|
31 |
scan_read = transformer(scan_dict)
|
32 |
+
inimg_raw = scan_read['image'].swapaxes(0, 2)
|
33 |
|
34 |
tvolslices, xnew_box = lungmask.utils.preprocess(inimg_raw, resolution=[256, 256])
|
35 |
tvolslices[tvolslices > 600] = 600
|
|
|
42 |
timage_res = np.empty((np.append(0, tvolslices[0].shape)), dtype=np.uint8)
|
43 |
|
44 |
with torch.no_grad():
|
45 |
+
for X in tqdm(dataloader_val):
|
46 |
X = X.float().to(device)
|
47 |
prediction = model(X)
|
48 |
pls = torch.max(prediction, 1)[1].detach().cpu().numpy().astype(np.uint8)
|
|
|
58 |
outmask = np.swapaxes(outmask, 0, 2)
|
59 |
#outmask = np.flip(outmask, 0)
|
60 |
|
|
|
61 |
return outmask.astype(np.uint8), scan_read['image_meta_dict']['affine']
|
62 |
|
63 |
def calculate_extremes(image, annotation_value):
|
|
|
143 |
preprocess_dump['pixdim'] = im['image_meta_dict']['pixdim'][1:4]
|
144 |
preprocess_dump['org_affine'] = im['image_meta_dict']['affine']
|
145 |
|
146 |
+
print("Segmenting lungs...")
|
147 |
masked_lungs = mask_lung(image_path, 5)
|
148 |
right_lung_extreme = calculate_extremes(masked_lungs[0], 1)
|
149 |
preprocess_dump['right_extremes'] = right_lung_extreme
|
|
|
250 |
|
251 |
stitched = np.logical_or(left, right).astype(int)
|
252 |
|
253 |
+
return stitched
|
lungtumormask/mask.py
CHANGED
@@ -32,4 +32,4 @@ def mask(image_path, save_path):
|
|
32 |
|
33 |
print(f"Storing segmentation at {save_path}")
|
34 |
nimage = nibabel.Nifti1Image(infered, preprocess_dump['org_affine'])
|
35 |
-
nibabel.save(nimage, save_path)
|
|
|
32 |
|
33 |
print(f"Storing segmentation at {save_path}")
|
34 |
nimage = nibabel.Nifti1Image(infered, preprocess_dump['org_affine'])
|
35 |
+
nibabel.save(nimage, save_path)
|
lungtumormask/network.py
CHANGED
@@ -148,4 +148,4 @@ class UNet_double(nn.Module):
|
|
148 |
return conv
|
149 |
|
150 |
def forward(self, x, box = None) -> torch.Tensor:
|
151 |
-
return self.activation(self.model1(x))
|
|
|
148 |
return conv
|
149 |
|
150 |
def forward(self, x, box = None) -> torch.Tensor:
|
151 |
+
return self.activation(self.model1(x))
|