added batch mode. Can now provide directory of CTs/.nii files to be predicted on. Also refactored build workflow

.github/workflows/build.yaml

@@ -34,10 +34,7 @@ jobs:
         run: pip install wheel setuptools
 
       - name: Build wheel
-        run: |
-          python setup.py bdist_wheel --universal
-          cd dist
-          ls
+        run: python setup.py bdist_wheel --universal
 
       - name: Install program for ${{matrix.TARGET}}
         run: ${{matrix.CMD_BUILD}}

livermask/livermask.py

@@ -12,6 +12,7 @@ import warnings
 import argparse
 import pkg_resources
 import tensorflow as tf
+import logging as log
 
 
 os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'  # due to this: https://github.com/tensorflow/tensorflow/issues/35029
@@ -38,7 +39,17 @@ def get_model(output):
     md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
     gdown.cached_download(url, output, md5=md5) #, postprocess=gdown.extractall)
 
-def func(path, output, cpu):
+def verboseHandler(verbose):
+    if verbose:
+        log.basicConfig(format="%(levelname)s: %(message)s", level=log.DEBUG)
+        log.info("Verbose output.")
+    else:
+        log.basicConfig(format="%(levelname)s: %(message)s")
+
+def func(path, output, cpu, verbose):
+    # enable verbose or not
+    verboseHandler(verbose)
+
     cwd = "/".join(os.path.realpath(__file__).replace("\\", "/").split("/")[:-1]) + "/"
     name = cwd + "model.h5"
 
@@ -48,83 +59,97 @@ def func(path, output, cpu):
     # load model
     model = load_model(name, compile=False)
 
-    print("preprocessing...")
-    nib_volume = nib.load(path)
-    new_spacing = [1., 1., 1.]
-    resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
-    data = resampled_volume.get_data().astype('float32')
-
-    curr_shape = data.shape
-
-    # resize to get (512, 512) output images
-    img_size = 512
-    data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
-
-    # intensity normalization
-    intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
-    data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
-
-    # fix orientation
-    data = np.rot90(data, k=1, axes=(0, 1))
-    data = np.flip(data, axis=0)
-
-    print("predicting...")
-    # predict on data
-    pred = np.zeros_like(data).astype(np.float32)
-    for i in tqdm(range(data.shape[-1]), "pred: "):
-        pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
-    del data 
-
-    # threshold
-    pred = (pred >= 0.4).astype(int)
-
-    # fix orientation back
-    pred = np.flip(pred, axis=0)
-    pred = np.rot90(pred, k=-1, axes=(0, 1))
-
-    print("resize back...")
-    # resize back from 512x512
-    pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
-    pred = (pred >= 0.5).astype(np.float32)
-
-    print("morphological post-processing...")
-    # morpological post-processing
-    # 1) first erode
-    pred = binary_erosion(pred.astype(bool), ball(3)).astype(np.float32)
-
-    # 2) keep only largest connected component
-    labels = label(pred)
-    regions = regionprops(labels)
-    area_sizes = []
-    for region in regions:
-        area_sizes.append([region.label, region.area])
-    area_sizes = np.array(area_sizes)
-    tmp = np.zeros_like(pred)
-    tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
-    pred = tmp.copy()
-    del tmp, labels, regions, area_sizes
-
-    # 3) dilate
-    pred = binary_dilation(pred.astype(bool), ball(3))
-
-    # 4) remove small holes
-    pred = remove_small_holes(pred.astype(bool), area_threshold=0.001*np.prod(pred.shape)).astype(np.float32)
-
-    print("saving...")
-    pred = pred.astype(np.uint8)
-    img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
-    resampled_lab = resample_from_to(img, nib_volume, order=0)
-    nib.save(resampled_lab, output)
-
+    if not os.path.isdir(path):
+        paths = [path]
+    else:
+        paths = [path + "/" + p for p in os.listdir(path)]
+
+    multiple_flag = len(paths) > 1
+    if multiple_flag:
+        os.makedirs(output + "/", exist_ok=True)
+
+    for curr in tqdm(paths, "CT:"):
+        log.info("preprocessing...")
+        nib_volume = nib.load(curr)
+        new_spacing = [1., 1., 1.]
+        resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
+        data = resampled_volume.get_data().astype('float32')
+
+        curr_shape = data.shape
+
+        # resize to get (512, 512) output images
+        img_size = 512
+        data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
+
+        # intensity normalization
+        intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
+        data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
+
+        # fix orientation
+        data = np.rot90(data, k=1, axes=(0, 1))
+        data = np.flip(data, axis=0)
+
+        log.info("predicting...")
+        # predict on data
+        pred = np.zeros_like(data).astype(np.float32)
+        for i in tqdm(range(data.shape[-1]), "pred: ", disable=not verbose):
+            pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
+        del data 
+
+        # threshold
+        pred = (pred >= 0.4).astype(int)
+
+        # fix orientation back
+        pred = np.flip(pred, axis=0)
+        pred = np.rot90(pred, k=-1, axes=(0, 1))
+
+        log.info("resize back...")
+        # resize back from 512x512
+        pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
+        pred = (pred >= 0.5).astype(np.float32)
+
+        log.info("morphological post-processing...")
+        # morpological post-processing
+        # 1) first erode
+        pred = binary_erosion(pred.astype(bool), ball(3)).astype(np.float32)
+
+        # 2) keep only largest connected component
+        labels = label(pred)
+        regions = regionprops(labels)
+        area_sizes = []
+        for region in regions:
+            area_sizes.append([region.label, region.area])
+        area_sizes = np.array(area_sizes)
+        tmp = np.zeros_like(pred)
+        tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
+        pred = tmp.copy()
+        del tmp, labels, regions, area_sizes
+
+        # 3) dilate
+        pred = binary_dilation(pred.astype(bool), ball(3))
+
+        # 4) remove small holes
+        pred = remove_small_holes(pred.astype(bool), area_threshold=0.001*np.prod(pred.shape)).astype(np.float32)
+
+        log.info("saving...")
+        pred = pred.astype(np.uint8)
+        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")
+        else:
+            nib.save(resampled_lab, output + ".nii")
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--input', metavar='--i', type=str, nargs='?',
-                    help="set path of which image to use.")
+                    help="set path of which image(s) to use.")
     parser.add_argument('--output', metavar='--o', type=str, nargs='?',
                     help="set path to store the output.")
     parser.add_argument('--cpu', action='store_true',
                     help="force using the CPU even if a GPU is available.")
+    parser.add_argument('--verbose', action='store_true',
+                    help="enable verbose.")
     ret = parser.parse_args(sys.argv[1:]); print(ret)
 
     if ret.cpu:
@@ -148,10 +173,10 @@ def main():
         raise ValueError("Please, provide an input.")
     if ret.output is None:
         raise ValueError("Please, provide an output.")
-    if not ret.input.endswith(".nii"):
-        raise ValueError("Image provided is not in the supported '.nii' format.")
-    if not ret.output.endswith(".nii"):
-        raise ValueError("Output name set is not in the supported '.nii' format.")
+    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 not os.path.isdir(ret.output) or "." in ret.output.split("/")[-1]:
+        raise ValueError("Output path provided is not a directory or a name (remove *.nii format from name).")
 
     # fix paths
     ret.input = ret.input.replace("\\", "/")