Added option to output probability map from CLI

lungtumormask/__main__.py

@@ -27,6 +27,10 @@ def main():
     if args.cpu:
         os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
     
+    # check if chosen threshold is in accepted range
+    if not (((args.threshold >= 0.0) and (args.threshold <= 1.0)) or (args.threshold == -1)):
+        raise ValueError("Chosen threshold must be -1 or in range [0.0, 1.0], but was:", args.threshold)
+    
     # import method here to enable faster testing
     from lungtumormask import mask
     

lungtumormask/dataprocessing.py

@@ -226,20 +226,30 @@ def post_process(left, right, preprocess_dump, lung_filter, threshold, radius):
     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)
+    if threshold != -1:
+        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'])
 
-    stitched = np.logical_or(left, right).astype(int)
+    # fuse left and right lung preds together (but keep probability maps if available)
+    stitched = np.maximum(left, right)
+    del left, right
 
     # filter tumor predictions outside the predicted lung area
     if lung_filter:
         stitched[preprocess_dump['lungmask'] == 0] = 0
     
-    # final post-processing - fix fragmentation
-    for i in range(stitched.shape[-1]):
-        stitched[..., i] = binary_closing(stitched[..., i], footprint=disk(radius=radius))
+    # final post-processing - fix fragmentation (only relevant for binary volume)
+    if threshold != -1:
+        for i in range(stitched.shape[-1]):
+            stitched[..., i] = binary_closing(stitched[..., i], footprint=disk(radius=radius))
+
+    # for threshold != -1, set result to uint8 dtype, else float32 (for probability map)
+    if threshold == -1:
+        stitched = stitched.astype("float32")
+    else:
+        stitched = stitched.astype("uint8")
 
     return stitched

lungtumormask/mask.py

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