Enable getting probability map when running from Docker (#250)

* Get probabilities in RunCellpose from docker

include a conditional statement to get the right flow for probabilities when running RunCellpose from docker image

* add probability map rescaling option when running from Python

Author: emiglietta

active_plugins/runcellpose.py

@@ -97,7 +97,7 @@ class RunCellpose(ImageSegmentation):
 
     module_name = "RunCellpose"
 
-    variable_revision_number = 5
+    variable_revision_number = 6
 
     doi = {
         "Please cite the following when using RunCellPose:": "https://doi.org/10.1038/s41592-020-01018-x",
@@ -111,8 +111,8 @@ def create_settings(self):
             text="Rescale images before running Cellpose",
             value=True,
             doc="""\
-Reminds the user that the  normalization step will be performed to ensure suimilar segmentation behaviour inthe RunCellpose
-modul and the Cellpose app.
+Reminds the user that the  normalization step will be performed to ensure suimilar segmentation behaviour in the RunCellpose
+module and the Cellpose app.
 """
         )
 
@@ -346,6 +346,14 @@ def set_directory_fn(path):
             doc="""
 If you do not want to include any object masks that are not in full view in the image, you can have the masks that have pixels touching the the edges removed.
 The default is set to "Yes".
+""",
+        )
+
+        self.probability_rescale_setting = Binary(
+            text="Rescale probability map?",
+            value=True,
+            doc="""
+Activate to rescale probability map to 0-255 (which matches the scale used when running this module from Docker)
 """,
         )
 
@@ -375,6 +383,7 @@ def settings(self):
             self.omni,
             self.invert,
             self.remove_edge_masks,
+            self.probability_rescale_setting,
         ]
 
     def visible_settings(self):
@@ -416,6 +425,8 @@ def visible_settings(self):
 
         if self.save_probabilities.value:
             vis_settings += [self.probabilities_name]
+            if self.docker_or_python.value == 'Python':
+                vis_settings += [self.probability_rescale_setting]
 
         vis_settings += [self.use_averaging, self.use_gpu]
 
@@ -630,9 +641,18 @@ def run(self, workspace):
         objects.add_objects(y, y_name)
 
         if self.save_probabilities.value:
+            if self.docker_or_python.value == "Docker":
+                # get rid of extra dimension
+                prob_map = numpy.squeeze(flows[1], axis=0) # ranges 0-255
+            else:
+                prob_map = flows[2]
+                rescale_prob_map = prob_map.copy()
+                prob_map01 = numpy.percentile(rescale_prob_map, 1)
+                prob_map99 = numpy.percentile(rescale_prob_map, 99)
+                prob_map = numpy.clip((rescale_prob_map - prob_map01) / (prob_map99 - prob_map01), a_min=0, a_max=1)
             # Flows come out sized relative to CellPose's inbuilt model size.
             # We need to slightly resize to match the original image.
-            size_corrected = skimage.transform.resize(flows[2], y_data.shape)
+            size_corrected = skimage.transform.resize(prob_map, y_data.shape)
             prob_image = Image(
                 size_corrected,
                 parent_image=x.parent_image,
@@ -723,6 +743,9 @@ def upgrade_settings(self, setting_values, variable_revision_number, module_name
         if variable_revision_number == 4:
             setting_values = [setting_values[0]] + ['No'] + setting_values[1:]
             variable_revision_number = 5
+        if variable_revision_number == 5:
+            setting_values = setting_values + [False]
+            variable_revision_number = 6
         return setting_values, variable_revision_number