Merge pull request #10 from andreped/interpolation

Fixed interpolation artifact + added morphological post-processing [no ci]

.github/workflows/build.yml

@@ -39,7 +39,7 @@ jobs:
     strategy:
       matrix:
         os: [windows-2019, ubuntu-20.04, macos-10.15]
-        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10"]
+        python-version: ["3.7", "3.8", "3.9", "3.10"]
 
     steps:
       - uses: actions/checkout@v1
@@ -67,4 +67,4 @@ jobs:
         run: lungtumormask --help
 
       - name: Test inference
-        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --threshold 0.3 --lung-filter
+        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --threshold 0.3 --lung-filter --radius 3

.gitignore

@@ -1,3 +1,5 @@
 .idea/
 *.nii
-*.nii.gz
+*.nii.gz
+*__pycache__/
+*.egg-info

README.md

@@ -13,7 +13,7 @@ A pretrained model is made available in a command line tool and can be used as y
 
 ## [Installation](https://github.com/VemundFredriksen/LungTumorMask#installation)
 
-Software has been tested against Python `3.6-3.10`.
+Software has been tested against Python `3.7-3.10`.
 
 Stable latest release:
 ```

lungtumormask/__main__.py

@@ -14,10 +14,12 @@ def main():
     parser.add_argument('input', metavar='input', type=path, help='Path to the input image, should be .nifti')
     parser.add_argument('output', metavar='output', type=str, help='Filepath for output tumormask')
     parser.add_argument('--lung-filter', action='store_true', help='whether to apply lungmask postprocessing.')
-    parser.add_argument('--threshold', metavar='threshold', type=float, default=0.4, 
+    parser.add_argument('--threshold', metavar='threshold', type=float, default=0.5, 
                         help='which threshold to use for assigning voxel-wise classes.')
+    parser.add_argument('--radius', metavar='radius', type=int, default=5,
+                        help='which radius to use for morphological post-processing segmentation smoothing.')
 
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
 
-    mask.mask(args.input, args.output, args.lung_filter, args.threshold)
+    mask.mask(args.input, args.output, args.lung_filter, args.threshold, args.radius)

lungtumormask/dataprocessing.py

@@ -8,6 +8,7 @@ import torch
 import numpy as np
 from monai.transforms import (Compose, LoadImaged, ToNumpyd, ThresholdIntensityd, AddChanneld, NormalizeIntensityd, SpatialCropd, DivisiblePadd, Spacingd, SqueezeDimd)
 from tqdm import tqdm
+from skimage.morphology import binary_closing, ball
 
 def mask_lung(scan_path, batch_size=20):
     model = lungmask.mask.get_model('unet', 'R231')
@@ -50,7 +51,6 @@ def mask_lung(scan_path, batch_size=20):
 
     outmask = lungmask.utils.postprocessing(timage_res)
 
-
     outmask = np.asarray(
         [lungmask.utils.reshape_mask(outmask[i], xnew_box[i], inimg_raw.shape[1:]) for i in range(outmask.shape[0])],
         dtype=np.uint8)
@@ -61,7 +61,6 @@ def mask_lung(scan_path, batch_size=20):
     return outmask.astype(np.uint8), scan_read['image_meta_dict']['affine']
 
 def calculate_extremes(image, annotation_value):
-
     holder = np.copy(image)
 
     x_min = float('inf')
@@ -98,7 +97,6 @@ def calculate_extremes(image, annotation_value):
     return ((x_min, x_max), (y_min, y_max), (z_min, z_max))
 
 def process_lung_scan(scan_dict, extremes):
-
     load_transformer = Compose(
         [
             LoadImaged(keys=["image"]),
@@ -123,15 +121,12 @@ def process_lung_scan(scan_dict, extremes):
     )
 
     processed_2 = transformer_1(processed_1)
-
     affine = processed_2['image_meta_dict']['affine']
-
     normalized_image = processed_2['image']
 
     return normalized_image, affine
 
 def preprocess(image_path):
-
     preprocess_dump = {}
 
     scan_dict = {
@@ -208,7 +203,6 @@ def find_pad_edge(original):
 
     return a_min, a_max + 1, b_min, b_max + 1, c_min, c_max + 1
 
-
 def remove_pad(mask, original):
     a_min, a_max, b_min, b_max, c_min, c_max = find_pad_edge(original)
     
@@ -216,27 +210,25 @@ def remove_pad(mask, original):
 
 def voxel_space(image, target):
     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]
-    image = ThresholdIntensity(above = False, threshold = 0.5, cval = 1)(image)
-    image = ThresholdIntensity(above = True, threshold = 0.5, cval = 0)(image)
 
     return image
 
 def stitch(org_shape, cropped, roi):
-    holder = np.zeros(org_shape)
+    holder = np.zeros(org_shape, dtype="float32")
     holder[roi[0][0]:roi[0][1], roi[1][0]:roi[1][1], roi[2][0]:roi[2][1]] = cropped
 
     return holder
 
-def post_process(left_mask, right_mask, preprocess_dump, lung_filter, threshold):
-    left_mask = (left_mask >= threshold).astype(int)
-    right_mask = (right_mask >= threshold).astype(int)
-
-    left = remove_pad(left_mask, preprocess_dump['left_lung'].squeeze(0).squeeze(0).numpy())
-    right = remove_pad(right_mask, preprocess_dump['right_lung'].squeeze(0).squeeze(0).numpy())
+def post_process(left, right, preprocess_dump, lung_filter, threshold, radius):
+    left = remove_pad(left, preprocess_dump['left_lung'].squeeze(0).squeeze(0).numpy())
+    right = remove_pad(right, preprocess_dump['right_lung'].squeeze(0).squeeze(0).numpy())
 
     left = voxel_space(left, preprocess_dump['left_extremes'])
     right = voxel_space(right, preprocess_dump['right_extremes'])
 
+    left = (left >= threshold).astype(int)
+    right = (right >= threshold).astype(int)
+
     left = stitch(preprocess_dump['org_shape'], left, preprocess_dump['left_extremes'])
     right = stitch(preprocess_dump['org_shape'], right, preprocess_dump['right_extremes'])
 
@@ -245,5 +237,8 @@ def post_process(left_mask, right_mask, preprocess_dump, lung_filter, threshold)
     # filter tumor predictions outside the predicted lung area
     if lung_filter:
         stitched[preprocess_dump['lungmask'] == 0] = 0
+    
+    # final post-processing - fix fragmentation
+    stitched = binary_closing(stitched, footprint=ball(radius=radius))
 
     return stitched

lungtumormask/mask.py

@@ -15,7 +15,7 @@ def load_model():
     model.eval()
     return model
 
-def mask(image_path, save_path, lung_filter, threshold):
+def mask(image_path, save_path, lung_filter, threshold, radius):
     print("Loading model...")
     model = load_model()
 
@@ -27,7 +27,7 @@ def mask(image_path, save_path, lung_filter, threshold):
     right = model(preprocess_dump['right_lung']).squeeze(0).squeeze(0).detach().numpy()
 
     print("Post-processing image...")
-    inferred = post_process(left, right, preprocess_dump, lung_filter, threshold).astype("uint8")
+    inferred = post_process(left, right, preprocess_dump, lung_filter, threshold, radius).astype("uint8")
 
     print(f"Storing segmentation at {save_path}")
     nimage = nibabel.Nifti1Image(inferred, preprocess_dump['org_affine'])