Swapped 3D with 2D closing operation [no ci]

lungtumormask/__main__.py

@@ -16,7 +16,7 @@ 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,
+    parser.add_argument('--radius', metavar='radius', type=int, default=1,
                         help='which radius to use for morphological post-processing segmentation smoothing.')
 
     argsin = sys.argv[1:]

lungtumormask/dataprocessing.py

@@ -8,7 +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
+from skimage.morphology import binary_closing, disk
 
 def mask_lung(scan_path, batch_size=20):
     model = lungmask.mask.get_model('unet', 'R231')
@@ -239,6 +239,7 @@ def post_process(left, right, preprocess_dump, lung_filter, threshold, radius):
         stitched[preprocess_dump['lungmask'] == 0] = 0
     
     # final post-processing - fix fragmentation
-    stitched = binary_closing(stitched, footprint=ball(radius=radius))
+    for i in range(stitched.shape[-1]):
+        stitched[..., i] = binary_closing(stitched[..., i], footprint=disk(radius=radius))
 
     return stitched