Support for setting threshold through argparse

.github/workflows/build.yml

@@ -67,7 +67,4 @@ jobs:
         run: lungtumormask --help
 
       - name: Test inference
-        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz
-      
-      #- name: Test lungmask postprocessing
-      #  run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --lung-filter
+        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --threshold 0.3 --lung-filter

lungtumormask/__main__.py

@@ -13,10 +13,11 @@ 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_argument('--lung-filter', action='store_true',
-                       help="whether to apply lungmask postprocessing.")
+    parser.add_argument('--lung-filter', action='store_true', help='whether to apply lungmask postprocessing.')
+    parser.add_argument('--threshold', metavar='threshold', type=int, default=0.4, 
+                        help='which threshold to use for assigning voxel-wise classes.')
 
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
 
-    mask.mask(args.input, args.output, args.lung_filter)
+    mask.mask(args.input, args.output, args.lung_filter, args.threshold)

lungtumormask/dataprocessing.py

@@ -211,12 +211,14 @@ def find_pad_edge(original):
 
 def remove_pad(mask, original):
     a_min, a_max, b_min, b_max, c_min, c_max = find_pad_edge(original)
+    
     return mask[a_min:a_max, b_min:b_max, c_min: c_max]
 
 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):
@@ -225,9 +227,9 @@ def stitch(org_shape, cropped, roi):
 
     return holder
 
-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)
+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())

lungtumormask/mask.py

@@ -15,7 +15,7 @@ def load_model():
     model.eval()
     return model
 
-def mask(image_path, save_path, lung_filter):
+def mask(image_path, save_path, lung_filter, threshold):
     print("Loading model...")
     model = load_model()
 
@@ -27,7 +27,7 @@ def mask(image_path, save_path, lung_filter):
     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).astype("uint8")
+    inferred = post_process(left, right, preprocess_dump, lung_filter, threshold).astype("uint8")
 
     print(f"Storing segmentation at {save_path}")
     nimage = nibabel.Nifti1Image(inferred, preprocess_dump['org_affine'])