Support for setting closing smoothing radius [no ci]

lungtumormask/__main__.py

@@ -16,8 +16,10 @@ def main():
     parser.add_argument('--lung-filter', action='store_true', help='whether to apply lungmask postprocessing.')
     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

@@ -219,7 +219,7 @@ def stitch(org_shape, cropped, roi):
 
     return holder
 
-def post_process(left, right, preprocess_dump, lung_filter, threshold):
+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())
 
@@ -239,6 +239,6 @@ def post_process(left, right, preprocess_dump, lung_filter, threshold):
         stitched[preprocess_dump['lungmask'] == 0] = 0
     
     # final post-processing - fix fragmentation
-    stitched = binary_closing(stitched, footprint=ball(radius=5))
+    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'])