option to merge or split nuc mem

app.py

@@ -6,7 +6,6 @@ from numpy.typing import ArrayLike
 import numpy as np
 from skimage import exposure
 from skimage.transform import resize
-from skimage import img_as_float
 from skimage.util import invert
 import cmap
 
@@ -86,7 +85,7 @@ class VSGradio:
         nuc_rgb = apply_colormap(nuc_pred, green_colormap)
         mem_rgb = apply_colormap(mem_pred, magenta_colormap)
 
-        return nuc_rgb, mem_rgb
+        return nuc_rgb, mem_rgb  # Return both nucleus and membrane images
 
 
 def apply_colormap(prediction, colormap: cmap.Colormap):
@@ -103,6 +102,11 @@ def apply_colormap(prediction, colormap: cmap.Colormap):
     return rgb_image_uint8
 
 
+def merge_images(nuc_rgb: ArrayLike, mem_rgb: ArrayLike) -> ArrayLike:
+    """Merge nucleus and membrane images into a single RGB image."""
+    return np.maximum(nuc_rgb, mem_rgb)
+
+
 def apply_image_adjustments(image, invert_image: bool, gamma_factor: float):
     """Applies all the image adjustments (invert, contrast, gamma) in sequence"""
     # Apply invert
@@ -188,9 +192,12 @@ if __name__ == "__main__":
                 output_membrane = gr.Image(
                     type="numpy", image_mode="RGB", label="VS Membrane"
                 )
+                merged_image = gr.Image(
+                    type="numpy", image_mode="RGB", label="Merged Image", visible=False
+                )
 
         # Checkbox for applying invert
-        preprocess_invert = gr.Checkbox(label="Apply Invert", value=False)
+        preprocess_invert = gr.Checkbox(label="Invert Image", value=False)
 
         # Slider for gamma adjustment
         gamma_factor = gr.Slider(
@@ -204,6 +211,9 @@ if __name__ == "__main__":
             placeholder="Enter cell diameter in microns",
         )
 
+        # Checkbox for merging predictions
+        merge_checkbox = gr.Checkbox(label="Merge Predictions into one image", value=False)
+
         # Update the adjusted image based on all the transformations
         input_image.change(
             fn=apply_image_adjustments,
@@ -223,14 +233,37 @@ if __name__ == "__main__":
             outputs=adjusted_image,
         )
 
-        # Button to trigger prediction
+        # Button to trigger prediction and update the output images
         submit_button = gr.Button("Submit")
 
-        # Define what happens when the button is clicked (send adjusted image to predict)
+        # Function to handle prediction and merging if needed
+        def submit_and_merge(inp, cell_diameter, merge):
+            nucleus, membrane = vsgradio.predict(inp, cell_diameter)
+            if merge:
+                merged = merge_images(nucleus, membrane)
+                return merged, gr.update(visible=True), nucleus, gr.update(visible=False), membrane,gr.update(visible=False)
+            else:
+                return None, gr.update(visible=False), nucleus, gr.update(visible=True), membrane, gr.update(visible=True)
+
         submit_button.click(
-            vsgradio.predict,
-            inputs=[adjusted_image, cell_diameter],
-            outputs=[output_nucleus, output_membrane],
+            fn=submit_and_merge,
+            inputs=[adjusted_image, cell_diameter, merge_checkbox],
+            outputs=[merged_image, merged_image, output_nucleus, output_nucleus,output_membrane, output_membrane],
+        )
+
+        # Function to handle merging the two predictions after they are shown
+        def merge_predictions_fn(nucleus_image, membrane_image, merge):
+            if merge:
+                merged = merge_images(nucleus_image, membrane_image)
+                return merged, gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+            else:
+                return None, gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+
+        # Toggle between merged and separate views when the checkbox is checked
+        merge_checkbox.change(
+            fn=merge_predictions_fn,
+            inputs=[output_nucleus, output_membrane, merge_checkbox],
+            outputs=[merged_image, merged_image, output_nucleus, output_membrane],
         )
 
         # Example images and article