Added support to state whether to filter lungs or not - default no

lungtumormask/__main__.py

@@ -13,8 +13,10 @@ def main():
     parser = argparse.ArgumentParser()
     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_argment('lung-filter', metavar='post-process', action='store_true',
+                       help="whether to apply lungmask postprocessing.")
 
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
 
-    mask.mask(args.input, args.output)
+    mask.mask(args.input, args.output, args.lung_filter)

lungtumormask/dataprocessing.py

@@ -217,17 +217,15 @@ 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[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):
+def post_process(left_mask, right_mask, preprocess_dump, lung_filter):
     left_mask = (left_mask >= 0.5).astype(int)
     right_mask = (right_mask >= 0.5).astype(int)
 
@@ -243,6 +241,7 @@ def post_process(left_mask, right_mask, preprocess_dump):
     stitched = np.logical_or(left, right).astype(int)
 
     # filter tumor predictions outside the predicted lung area
-    stitched[preprocess_dump['lungmask'] == 0] = 0
+    if lung_filter:
+        stitched[preprocess_dump['lungmask'] == 0] = 0
 
     return stitched

lungtumormask/mask.py

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