Merge pull request #78 from JoHof/code_cleaning

Code cleaning

.flake8

@@ -0,0 +1,4 @@
+[flake8]
+max-line-length = 88
+ignore = E203, E266, E501, W503, F403, F401
+select = B,C,E,F,W,T4,B9

.gitignore

@@ -1,4 +1,6 @@
 __pycache__/
 .idea/
-lungmask.egg-info/
-lungmask/.ipynb_checkpoints/
+*.egg-info
+**/.ipynb_checkpoints/
+*.pyc
+**/.coverage*

.pre-commit-config.yaml

@@ -0,0 +1,42 @@
+repos:
+
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+    -   id: check-ast
+    -   id: debug-statements
+    -   id: end-of-file-fixer
+        exclude: resources/
+    -   id: requirements-txt-fixer
+    -   id: trailing-whitespace
+    -   id: check-added-large-files
+        args: ['--maxkb=500']
+
+
+-   repo: https://github.com/psf/black
+    rev: 23.3.0
+    hooks:
+    -   id: black
+        language_version: python3
+
+-   repo: https://github.com/pycqa/flake8
+    rev: 6.0.0
+    hooks:
+    - id: flake8
+
+-   repo: https://github.com/pre-commit/mirrors-isort
+    rev: v5.10.1
+    hooks:
+    -   id: isort
+
+-   repo: https://github.com/jorisroovers/gitlint
+    rev: v0.19.1
+    hooks:
+    - id: gitlint
+      entry: gitlint -c title-max-length.line-length=60 --ignore B6,B7 --msg-filename
+      stages: [commit-msg]
+
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.2.0
+    hooks:
+      - id: mypy

README.md

@@ -4,7 +4,7 @@ This package provides trained U-net models for lung segmentation. For now, four
 
 - U-net(R231): This model was trained on a large and diverse dataset that covers a wide range of visual variabiliy. The model performs segmentation on individual slices, extracts right-left lung seperately includes airpockets, tumors and effusions. The trachea will not be included in the lung segmentation. https://doi.org/10.1186/s41747-020-00173-2
 
-- U-net(LTRCLobes): This model was trained on a subset of the [LTRC](https://ltrcpublic.com) dataset. The model performs segmentation of individual lung-lobes but yields limited performance when dense pathologies are present or when fissures are not visible at every slice. 
+- U-net(LTRCLobes): This model was trained on a subset of the [LTRC](https://ltrcpublic.com) dataset. The model performs segmentation of individual lung-lobes but yields limited performance when dense pathologies are present or when fissures are not visible at every slice.
 
 - U-net(LTRCLobes_R231): This will run the R231 and LTRCLobes model and fuse the results. False negatives from LTRCLobes will be filled by R231 predictions and mapped to a neighbor label. False positives from LTRCLobes will be removed. The fusing process is computationally intensive and can, depdending on the data and results, take up to several minutes per volume.
 
@@ -60,33 +60,34 @@ lungmask -h
 ### As a python module:
 
 ```
-from lungmask import mask
+from lungmask import LMInferer
 import SimpleITK as sitk
 
+inferer = LMInferer()
+
 input_image = sitk.ReadImage(INPUT)
-segmentation = mask.apply(input_image)  # default model is U-net(R231)
+segmentation = inferer.apply(input_image)  # default model is U-net(R231)
 ```
 input_image has to be a SimpleITK object.
 
 Load an alternative model like so:
 ```
-model = mask.get_model('unet','LTRCLobes')
-segmentation = mask.apply(input_image, model)
+inferer = LMInferer(modelname="R231CovidWeb")
 ```
 
-To use the model fusing capability for LTRCLobes_R231 use:
+To use the model fusing capability for (e.g. LTRCLobes_R231) use:
 ```
-segmentation = mask.apply_fused(input_image)
+inferer = LMInferer(modelname='LTRCLobes', fillmodel='R231')
 ```
 
 #### Numpy array support
 As of version 0.2.9, numpy arrays are supported as input volumes. This mode assumes the input numpy array has the following format for each axis:
 * first axis containing slices
 * second axis with chest to back
-* third axis with right to left 
+* third axis with right to left
 
 ## Limitations
-The model works on full slices only. The slice to process has to show the full lung and the lung has to be surrounded by tissue in order to get segmented. However, the model is quite stable to cases with a cropped field of view as long as the lung is surrounded by tissue. 
+The model works on full slices only. The slice to process has to show the full lung and the lung has to be surrounded by tissue in order to get segmented. However, the model is quite stable to cases with a cropped field of view as long as the lung is surrounded by tissue.
 
 ## COVID-19 Web
 ```

lungmask/__init__.py

@@ -0,0 +1 @@
+from .mask import LMInferer

lungmask/__main__.py

@@ -1,68 +1,136 @@
-import sys
 import argparse
 import logging
-from lungmask import mask
-from lungmask import utils
 import os
-import SimpleITK as sitk
-import pkg_resources
+import sys
+
 import numpy as np
+import pkg_resources  # type: ignore
+import SimpleITK as sitk
+
+from lungmask import mask, utils
 
 
 def path(string):
     if os.path.exists(string):
         return string
     else:
-        sys.exit(f'File not found: {string}')
+        sys.exit(f"File not found: {string}")
 
 
 def main():
     version = pkg_resources.require("lungmask")[0].version
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument('input', metavar='input', type=path, help='Path to the input image, can be a folder for dicoms')
-    parser.add_argument('output', metavar='output', type=str, help='Filepath for output lungmask')
-    parser.add_argument('--modeltype', help='Default: unet', type=str, choices=['unet'], default='unet')
-    parser.add_argument('--modelname', help="spcifies the trained model, Default: R231", type=str, choices=['R231','LTRCLobes','LTRCLobes_R231','R231CovidWeb'], default='R231')
-    parser.add_argument('--modelpath', help="spcifies the path to the trained model", default=None)
-    parser.add_argument('--classes', help="spcifies the number of output classes of the model", default=3)
-    parser.add_argument('--cpu', help="Force using the CPU even when a GPU is available, will override batchsize to 1", action='store_true')
-    parser.add_argument('--nopostprocess', help="Deactivates postprocessing (removal of unconnected components and hole filling", action='store_true')
-    parser.add_argument('--noHU', help="For processing of images that are not encoded in hounsfield units (HU). E.g. png or jpg images from the web. Be aware, results may be substantially worse on these images", action='store_true')
-    parser.add_argument('--batchsize', type=int, help="Number of slices processed simultaneously. Lower number requires less memory but may be slower.", default=20)
-    parser.add_argument('--version', help="Shows the current version of lungmask", action='version', version=version)
+
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "input",
+        metavar="input",
+        type=path,
+        help="Path to the input image, can be a folder for dicoms",
+    )
+    parser.add_argument(
+        "output", metavar="output", type=str, help="Filepath for output lungmask"
+    )
+    parser.add_argument(
+        "--modeltype", help="Default: unet", type=str, choices=["unet"], default="unet"
+    )
+    parser.add_argument(
+        "--modelname",
+        help="spcifies the trained model, Default: R231",
+        type=str,
+        choices=["R231", "LTRCLobes", "LTRCLobes_R231", "R231CovidWeb"],
+        default="R231",
+    )
+    parser.add_argument(
+        "--modelpath", help="spcifies the path to the trained model", default=None
+    )
+    parser.add_argument(
+        "--classes",
+        help="spcifies the number of output classes of the model",
+        default=3,
+    )
+    parser.add_argument(
+        "--cpu",
+        help="Force using the CPU even when a GPU is available, will override batchsize to 1",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--nopostprocess",
+        help="Deactivates postprocessing (removal of unconnected components and hole filling)",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--noHU",
+        help="For processing of images that are not encoded in hounsfield units (HU). E.g. png or jpg images from the web. Be aware, results may be substantially worse on these images",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--batchsize",
+        type=int,
+        help="Number of slices processed simultaneously. Lower number requires less memory but may be slower.",
+        default=20,
+    )
+    parser.add_argument(
+        "--noprogress",
+        action="store_true",
+        help="If set, no tqdm progress bar will be shown",
+    )
+    parser.add_argument(
+        "--version",
+        help="Shows the current version of lungmask",
+        action="version",
+        version=version,
+    )
 
     argsin = sys.argv[1:]
     args = parser.parse_args(argsin)
-    
+
     batchsize = args.batchsize
     if args.cpu:
         batchsize = 1
 
-    logging.info(f'Load model')
-
-    input_image = utils.get_input_image(args.input)
-    logging.info(f'Infer lungmask')
-    if args.modelname == 'LTRCLobes_R231':
-        assert args.modelpath is None, "Modelpath can not be specified for LTRCLobes_R231 mode"
-        result = mask.apply_fused(input_image, force_cpu=args.cpu, batch_size=batchsize, volume_postprocessing=not(args.nopostprocess), noHU=args.noHU)
+    logging.info("Load model")
+
+    input_image = utils.load_input_image(args.input, disable_tqdm=args.noprogress)
+    logging.info("Infer lungmask")
+    if args.modelname == "LTRCLobes_R231":
+        assert (
+            args.modelpath is None
+        ), "Modelpath can not be specified for LTRCLobes_R231 mode"
+        result = mask.apply_fused(
+            input_image,
+            force_cpu=args.cpu,
+            batch_size=batchsize,
+            volume_postprocessing=not (args.nopostprocess),
+            noHU=args.noHU,
+            tqdm_disable=args.noprogress,
+        )
     else:
-        model = mask.get_model(args.modeltype, args.modelname, args.modelpath, args.classes)
-        result = mask.apply(input_image, model, force_cpu=args.cpu, batch_size=batchsize, volume_postprocessing=not(args.nopostprocess), noHU=args.noHU)
-
+        model = mask.get_model(args.modelname, args.modelpath, args.classes)
+        result = mask.apply(
+            input_image,
+            model,
+            force_cpu=args.cpu,
+            batch_size=batchsize,
+            volume_postprocessing=not (args.nopostprocess),
+            noHU=args.noHU,
+            tqdm_disable=args.noprogress,
+        )
+
     if args.noHU:
-        file_ending = args.output.split('.')[-1]
+        file_ending = args.output.split(".")[-1]
         print(file_ending)
-        if file_ending in ['jpg','jpeg','png']:
-            result = (result/(result.max())*255).astype(np.uint8)
+        if file_ending in ["jpg", "jpeg", "png"]:
+            result = (result / (result.max()) * 255).astype(np.uint8)
         result = result[0]
-             
-    result_out= sitk.GetImageFromArray(result)
+
+    result_out = sitk.GetImageFromArray(result)
     result_out.CopyInformation(input_image)
-    logging.info(f'Save result to: {args.output}')
-    sys.exit(sitk.WriteImage(result_out, args.output))
+    logging.info(f"Save result to: {args.output}")
+    sitk.WriteImage(result_out, args.output)
 
 
 if __name__ == "__main__":
-    print('called as script')
+    print("called as script")
     main()