Merge pull request #18 from andreped/lungmask-fix

Fixed lungmask filtering + set batch size through args + bump v1.2.1 [no ci]

.github/workflows/build.yml

@@ -68,4 +68,4 @@ jobs:
         run: lungtumormask --help
 
       - name: Test inference
-        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --threshold 0.3 --lung-filter --radius 3
+        run: lungtumormask samples/lung_001.nii.gz mask_001.nii.gz --threshold 0.3 --lung-filter --radius 3 --batch-size 8

.gitignore

@@ -3,3 +3,5 @@
 *.nii.gz
 *__pycache__/
 *.egg-info
+dist/
+build/

lungtumormask/__main__.py

@@ -1,7 +1,6 @@
 import sys
 import argparse
 import os
-from lungtumormask import mask
 
 def path(string):
     if os.path.exists(string):
@@ -18,8 +17,13 @@ def main():
                         help='which threshold to use for assigning voxel-wise classes.')
     parser.add_argument('--radius', metavar='radius', type=int, default=1,
                         help='which radius to use for morphological post-processing segmentation smoothing.')
+    parser.add_argument('--batch-size', metavar='batch-size', type=int, default=5,
+                        help='which batch size to use for lungmask inference.')
 
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
-
-    mask.mask(args.input, args.output, args.lung_filter, args.threshold, args.radius)
+    
+    # import method here to enable faster testing
+    from lungtumormask import mask
+    
+    mask.mask(args.input, args.output, args.lung_filter, args.threshold, args.radius, args.batch_size)

lungtumormask/dataprocessing.py

@@ -126,7 +126,7 @@ def process_lung_scan(scan_dict, extremes):
 
     return normalized_image, affine
 
-def preprocess(image_path):
+def preprocess(image_path, batch_size):
     preprocess_dump = {}
 
     scan_dict = {
@@ -139,8 +139,8 @@ def preprocess(image_path):
     preprocess_dump['org_affine'] = im['image_meta_dict']['affine']
 
     print("Segmenting lungs...")
-    masked_lungs = mask_lung(image_path, 5)
-    preprocess_dump['lungmask'] = masked_lungs
+    masked_lungs = mask_lung(image_path, batch_size=batch_size)
+    preprocess_dump['lungmask'] = masked_lungs[0]  # first output is binary segmentation of lungs
     right_lung_extreme = calculate_extremes(masked_lungs[0], 1)
     preprocess_dump['right_extremes'] = right_lung_extreme
     right_lung_processed = process_lung_scan(scan_dict, right_lung_extreme)

lungtumormask/mask.py

@@ -15,12 +15,12 @@ def load_model():
     model.eval()
     return model
 
-def mask(image_path, save_path, lung_filter, threshold, radius):
+def mask(image_path, save_path, lung_filter, threshold, radius, batch_size):
     print("Loading model...")
     model = load_model()
 
     print("Preprocessing image...")
-    preprocess_dump = preprocess(image_path)
+    preprocess_dump = preprocess(image_path, batch_size)
 
     print("Looking for tumors...")
     left = model(preprocess_dump['left_lung']).squeeze(0).squeeze(0).detach().numpy()

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
     name="lungtumormask",
     packages=find_packages(),
-    version='1.2.0',
+    version='1.2.1',
     author="Svein Ole M Sevle, Vemund Fredriksen, and André Pedersen",
     url="https://github.com/VemundFredriksen/LungTumorMask",
     license="MIT",