Update runstardist.py (#205)

* Update runstardist.py

Potentially resolves the issue#196 stardist plugin does not free gpu memory. 
Now you can runStardist and runCellpose on the same pipeline.

Tested on windows, need to test on Macs

* Only run if cuda is available

* Add optional GPU memory control code

---------

Co-authored-by: Beth Cimini <bethac07@users.noreply.github.com>

Author: mccruz07

active_plugins/runstardist.py

@@ -8,6 +8,8 @@
 from skimage.transform import resize
 from stardist.models import StarDist2D, StarDist3D
 from csbdeep.utils import normalize
+from numba import cuda
+import tensorflow as tf
 
 #################################
 #
@@ -206,6 +208,25 @@ def create_settings(self):
             maxval=1.0,
             doc=f"""\
 Prevent overlapping 
+""",
+        )
+
+        self.manage_gpu = Binary(
+            text="Manually set how much GPU memory each worker can use?",
+            value=False,
+            doc="""
+If enabled, you can manually set how much of the GPU memory each worker can use.
+This is likely to provide the most benefit on Macs. Do not use in a multi-GPU setup.""",
+        )
+
+        self.manual_GPU_memory_GB = Float(
+            text="GPU memory (in GB) for each worker",
+            value=0.5,
+            minval=0.0000001,
+            maxval=30,
+            doc="""\
+GPU memory in GB available to each worker. Value should be set such that this number times the number
+of workers in each copy of CellProfiler times the number of copies of CellProfiler running (if applicable) is <1
 """,
         )
 
@@ -224,6 +245,8 @@ def settings(self):
             self.model_choice3D,
             self.prob_thresh,
             self.nms_thresh,
+            self.manage_gpu,
+            self.manual_GPU_memory_GB,
         ]
 
     def visible_settings(self):
@@ -250,7 +273,10 @@ def visible_settings(self):
         if self.tile_image.value:
             vis_settings += [self.n_tiles_x, self.n_tiles_y]
 
-        vis_settings += [self.prob_thresh, self.nms_thresh, self.gpu_test]
+        vis_settings += [self.prob_thresh, self.nms_thresh, self.gpu_test, self.manage_gpu]
+
+        if self.manage_gpu.value:
+            vis_settings += [self.manual_GPU_memory_GB]
 
         return vis_settings
 
@@ -271,6 +297,23 @@ def run(self, workspace):
             raise ValueError(
                 "Greyscale images are not supported by this model. Please provide a color overlay."
             )
+        
+        # Stolen nearly wholesale from https://wiki.ncsa.illinois.edu/display/ISL20/Managing+GPU+memory+when+using+Tensorflow+and+Pytorch
+        if self.manage_gpu.value:
+            # First, Get a list of GPU devices
+            gpus = tf.config.list_physical_devices('GPU')
+            if len(gpus) > 0:            
+                # Restrict to only the first GPU.
+                tf.config.set_visible_devices(gpus[:1], device_type='GPU')
+                # Create a LogicalDevice with the appropriate memory limit
+                log_dev_conf = tf.config.LogicalDeviceConfiguration(
+                    memory_limit=self.manual_GPU_memory_GB.value*1024 # 2 GB
+                )
+                # Apply the logical device configuration to the first GPU
+                tf.config.set_logical_device_configuration(
+                    gpus[0],
+                    [log_dev_conf])
+
 
         if self.model.value == CUSTOM_MODEL:
             model_directory, model_name = os.path.split(
@@ -347,6 +390,8 @@ def run(self, workspace):
         objects.add_objects(y, self.y_name.value)
 
         self.add_measurements(workspace)
+        if cuda.is_available():
+            cuda.current_context().memory_manager.deallocations.clear()
 
         if self.show_window:
             workspace.display_data.x_data = x_data
@@ -389,11 +434,9 @@ def display(self, workspace, figure):
             )
 
     def do_check_gpu(self):
-        import tensorflow
-
-        if len(tensorflow.config.list_physical_devices("GPU")) > 0:
+        if len(tf.config.list_physical_devices("GPU")) > 0:
             message = "GPU appears to be working correctly!"
-            print("GPUs:", tensorflow.config.list_physical_devices("GPU"))
+            print("GPUs:", tf.config.list_physical_devices("GPU"))
         else:
             message = (
                 "GPU test failed. There may be something wrong with your configuration."