Merge pull request #20 from jpdefrutos/master

livermask/livermask.py

@@ -29,7 +29,7 @@ os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'  # due to this: https://github.
 warnings.filterwarnings('ignore', '.*output shape of zoom.*')  # mute some warnings
 
 
-def func(path, output, cpu, verbose, vessels):
+def func(path, output, cpu, verbose, vessels, extension):
     # enable verbose or not
     log = verboseHandler(verbose)
 
@@ -54,15 +54,15 @@ def func(path, output, cpu, verbose, vessels):
 
     for curr in tqdm(paths, "CT:"):
         # check if current file is a nifti file, if not, skip
-        if not curr.endswith(".nii"):
+        if not curr.endswith(".nii") or not curr.endswith(".nii.gz"):
             continue
 
         # perform liver parenchyma segmentation, launch it in separate process to properly clear memory
-        pred = liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name)
+        pred = liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name, extension)
 
         if vessels:
             # perform liver vessel segmentation
-            vessel_segmenter(curr, output, cpu, verbose, multiple_flag, pred, name_vessel)
+            vessel_segmenter(curr, output, cpu, verbose, multiple_flag, pred, name_vessel, extension)
 
 
 def main():
@@ -77,6 +77,8 @@ def main():
                     help="enable verbose.")
     parser.add_argument('--vessels', action='store_true',
                     help="segment vessels.")
+    parser.add_argument('--extension', metavar='--e', type=str, default=".nii",
+                    help="define the output extension. (default: .nii)")
     ret = parser.parse_args(sys.argv[1:]); print(ret)
 
     if ret.cpu:
@@ -105,10 +107,12 @@ def main():
     ret.input = ret.input.replace("\\", "/")
     ret.output = ret.output.replace("\\", "/")
 
-    if not os.path.isdir(ret.input) and not ret.input.endswith(".nii"):
-        raise ValueError("Input path provided is not in the supported '.nii' format or a directory.")
-    if ret.output.endswith(".nii") or "." in ret.output.split("/")[-1]:
+    if not os.path.isdir(ret.input) and not ret.input.endswith(".nii") and not ret.input.endswith(".nii.gz"):
+        raise ValueError("Input path provided is not in the supported '.nii' or '.nii.gz' formats or a directory.")
+    if ret.output.endswith(".nii") or ret.output.endswith(".nii.gz") or "." in ret.output.split("/")[-1]:
         raise ValueError("Output path provided is not a directory or a name (remove *.nii format from name).")
+    if ret.extension not in [".nii", ".nii.gz"]:
+        raise ValueError("Extension not supported. Expected: .nii or .nii.gz")
 
     func(*vars(ret).values())
 

livermask/utils/process.py

@@ -37,11 +37,11 @@ def intensity_normalization(volume, intensity_clipping_range):
     return result
 
 
-def liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name):
+def liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name, extension):
     # run inference in a different process
     mp.set_start_method('spawn', force=True)
     with mp.Pool(processes=1, maxtasksperchild=1) as p:  # , initializer=initializer)
-        result = p.map_async(liver_segmenter, ((curr, output, cpu, verbose, multiple_flag, name),))
+        result = p.map_async(liver_segmenter, ((curr, output, cpu, verbose, multiple_flag, name, extension),))
         log.info("getting result from process...")
         ret = result.get()[0]
     return ret
@@ -49,7 +49,7 @@ def liver_segmenter_wrapper(curr, output, cpu, verbose, multiple_flag, name):
 
 def liver_segmenter(params):
     try:
-        curr, output, cpu, verbose, multiple_flag, name = params
+        curr, output, cpu, verbose, multiple_flag, name, extension = params
 
         # load model
         model = load_model(name, compile=False)
@@ -123,9 +123,9 @@ def liver_segmenter(params):
         img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
         resampled_lab = resample_from_to(img, nib_volume, order=0)
         if multiple_flag:
-            nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-livermask.nii")
+            nib.save(resampled_lab, output + "/" + curr.split("/")[-1].split(".")[0] + "-livermask" + extension)
         else:
-            nib.save(resampled_lab, output + "-livermask.nii")
+            nib.save(resampled_lab, output + "-livermask" + extension)
 
         return pred
     except KeyboardInterrupt: