Disable GPU for app + moved CPU disable logic to inside run_analysis method

demo/app.py

@@ -29,7 +29,7 @@ def nifti_to_glb(path):
 def run_model(input_path):
     from livermask.utils.run import run_analysis
     
-    run_analysis(cpu=False, extension='.nii', path=input_path, output='prediction', verbose=True, vessels=False, name="./model.h5")
+    run_analysis(cpu=True, extension='.nii', path=input_path, output='prediction', verbose=True, vessels=False, name="./model.h5")
     
     #cmd_docker = ["python3", "-m", "livermask.livermask", "--input", input_path, "--output", "prediction", "--verbose"]
     #sp.check_call(cmd_docker, shell=True)  # @FIXME: shell=True here is not optimal -> starts a shell after calling script

livermask/livermask.py

@@ -25,26 +25,6 @@ def main():
     ret = parser.parse_args(sys.argv[1:])
     print(ret)
 
-    # only now do we call tensorflow, if necessary (to avoid redundant imports for livermask -h call)
-    import tensorflow as tf
-
-    if ret.cpu:
-        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-    if not tf.test.is_gpu_available():
-        tf.config.set_visible_devices([], 'GPU')
-        visible_devices = tf.config.get_visible_devices()
-    else:
-        gpus = tf.config.experimental.list_physical_devices('GPU')
-        try:
-            # Currently, memory growth needs to be the same across GPUs
-            for gpu in gpus:
-                tf.config.experimental.set_memory_growth(gpu, enable=True)
-            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
-            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
-        except RuntimeError as e:
-            # Memory growth must be set before GPUs have been initialized
-            print(e)
-
     if ret.input is None:
         raise ValueError("Please, provide an input.")
     if ret.output is None:

livermask/utils/run.py

@@ -29,6 +29,23 @@ def run_analysis(path, output, cpu, verbose, vessels, extension, name=None, name
     path = path.replace("\\", "/")
     output = output.replace("\\", "/")
 
+    if cpu:
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+    if not tf.test.is_gpu_available():
+        tf.config.set_visible_devices([], 'GPU')
+        visible_devices = tf.config.get_visible_devices()
+    else:
+        gpus = tf.config.experimental.list_physical_devices('GPU')
+        try:
+            # Currently, memory growth needs to be the same across GPUs
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, enable=True)
+            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+        except RuntimeError as e:
+            # Memory growth must be set before GPUs have been initialized
+            print(e)
+
     # enable verbose or not
     log = verboseHandler(verbose)