image adjustments, colormap and cell diameter

app.py

@@ -5,6 +5,10 @@ from huggingface_hub import hf_hub_download
 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
 
 
 class VSGradio:
@@ -31,9 +35,26 @@ class VSGradio:
         std = np.std(input)
         return (input - mean) / std
 
-    def predict(self, inp):
+    def preprocess_image_standard(self, input: ArrayLike):
+        # Perform standard preprocessing here
+        input = exposure.equalize_adapthist(input)
+        return input
+
+    def downscale_image(self, inp: ArrayLike, scale_factor: float):
+        """Downscales the image by the given scaling factor"""
+        height, width = inp.shape
+        new_height = int(height * scale_factor)
+        new_width = int(width * scale_factor)
+        return resize(inp, (new_height, new_width), anti_aliasing=True)
+
+    def predict(self, inp, cell_diameter: float):
         # Normalize the input and convert to tensor
         inp = self.normalize_fov(inp)
+        original_shape = inp.shape
+        # Resize the input image to the expected cell diameter
+        inp = apply_rescale_image(inp, cell_diameter, expected_cell_diameter=30)
+
+        # Convert the input to a tensor
         inp = torch.from_numpy(np.array(inp).astype(np.float32))
 
         # Prepare the input dictionary and move input to the correct device (GPU or CPU)
@@ -52,10 +73,60 @@ class VSGradio:
         # Post-process the model output and rescale intensity
         nuc_pred = pred[0, 0, 0]
         mem_pred = pred[0, 1, 0]
-        nuc_pred = exposure.rescale_intensity(nuc_pred, out_range=(0, 1))
-        mem_pred = exposure.rescale_intensity(mem_pred, out_range=(0, 1))
 
-        return nuc_pred, mem_pred
+        # Resize predictions back to the original image size
+        nuc_pred = resize(nuc_pred, original_shape, anti_aliasing=True)
+        mem_pred = resize(mem_pred, original_shape, anti_aliasing=True)
+
+        # Define colormaps
+        green_colormap = cmap.Colormap("green")  # Nucleus: black to green
+        magenta_colormap = cmap.Colormap("magenta")
+
+        # Apply the colormap to the predictions
+        nuc_rgb = apply_colormap(nuc_pred, green_colormap)
+        mem_rgb = apply_colormap(mem_pred, magenta_colormap)
+
+        return nuc_rgb, mem_rgb
+
+
+def apply_colormap(prediction, colormap: cmap.Colormap):
+    """Apply a colormap to a single-channel prediction image."""
+    # Ensure the prediction is within the valid range [0, 1]
+    prediction = exposure.rescale_intensity(prediction, out_range=(0, 1))
+
+    # Apply the colormap to get an RGB image
+    rgb_image = colormap(prediction)
+
+    # Convert the output from [0, 1] to [0, 255] for display
+    rgb_image_uint8 = (rgb_image * 255).astype(np.uint8)
+
+    return rgb_image_uint8
+
+
+def apply_image_adjustments(image, invert_image: bool, gamma_factor: float):
+    """Applies all the image adjustments (invert, contrast, gamma) in sequence"""
+    # Apply invert
+    if invert_image:
+        image = invert(image, signed_float=False)
+
+    # Apply gamma adjustment
+    image = exposure.adjust_gamma(image, gamma_factor)
+
+    return exposure.rescale_intensity(image, out_range=(0, 255)).astype(np.uint8)
+
+
+def apply_rescale_image(
+    image, cell_diameter: float, expected_cell_diameter: float = 30
+):
+    # Assume the model was trained with cells ~30 microns in diameter
+    # Resize the input image according to the scaling factor
+    scale_factor = expected_cell_diameter / float(cell_diameter)
+    image = resize(
+        image,
+        (int(image.shape[0] * scale_factor), int(image.shape[1] * scale_factor)),
+        anti_aliasing=True,
+    )
+    return image
 
 
 # Load the custom CSS from the file
@@ -64,7 +135,6 @@ def load_css(file_path):
         return file.read()
 
 
-# %%
 if __name__ == "__main__":
     # Download the model checkpoint from Hugging Face
     model_ckpt_path = hf_hub_download(
@@ -83,59 +153,102 @@ if __name__ == "__main__":
         "pretraining": False,
     }
 
+    vsgradio = VSGradio(model_config, model_ckpt_path)
+
     # Initialize the Gradio app using Blocks
     with gr.Blocks(css=load_css("style.css")) as demo:
         # Title and description
         gr.HTML(
             "<div class='title-block'>Image Translation (Virtual Staining) of cellular landmark organelles</div>"
         )
-        # Improved description block with better formatting
         gr.HTML(
             """
             <div class='description-block'>
                 <p><b>Model:</b> VSCyto2D</p>
-                <p>
-                    <b>Input:</b> label-free image (e.g., QPI or phase contrast) <br>
-                    <b>Output:</b> two virtually stained channels: one for the <b>nucleus</b> and one for the <b>cell membrane</b>.
-                </p>
-                <p>
-                    Check out our preprint: 
-                    <a href='https://www.biorxiv.org/content/10.1101/2024.05.31.596901' target='_blank'><i>Liu et al.,Robust virtual staining of landmark organelles</i></a>
-                </p>
+                <p><b>Input:</b> label-free image (e.g., QPI or phase contrast).</p>
+                <p><b>Output:</b> Virtual staining of nucleus and membrane.</p>
+                <p><b>Note:</b> The model works well with QPI, and sometimes generalizes to phase contrast and DIC. We continue to diagnose and improve generalization<p>
+                <p>Check out our preprint: <a href='https://www.biorxiv.org/content/10.1101/2024.05.31.596901' target='_blank'><i>Liu et al., Robust virtual staining of landmark organelles</i></a></p>
+                <p> For training, inference and evaluation of the model refer to the <a href='https://github.com/mehta-lab/VisCy/tree/main/examples/virtual_staining/dlmbl_exercise' target='_blank'>GitHub repository</a>.</p>
             </div>
             """
         )
 
-        vsgradio = VSGradio(model_config, model_ckpt_path)
-
         # Layout for input and output images
         with gr.Row():
             input_image = gr.Image(type="numpy", image_mode="L", label="Upload Image")
+            adjusted_image = gr.Image(
+                type="numpy", image_mode="L", label="Adjusted Image (Preview)"
+            )
+
             with gr.Column():
-                output_nucleus = gr.Image(type="numpy", label="VS Nucleus")
-                output_membrane = gr.Image(type="numpy", label="VS Membrane")
+                output_nucleus = gr.Image(
+                    type="numpy", image_mode="RGB", label="VS Nucleus"
+                )
+                output_membrane = gr.Image(
+                    type="numpy", image_mode="RGB", label="VS Membrane"
+                )
+
+        # Checkbox for applying invert
+        preprocess_invert = gr.Checkbox(label="Apply Invert", value=False)
+
+        # Slider for gamma adjustment
+        gamma_factor = gr.Slider(
+            label="Adjust Gamma", minimum=0.1, maximum=5.0, value=1.0, step=0.1
+        )
+
+        # Input field for the cell diameter in microns
+        cell_diameter = gr.Textbox(
+            label="Cell Diameter [um]",
+            value="30.0",
+            placeholder="Enter cell diameter in microns",
+        )
+
+        # Update the adjusted image based on all the transformations
+        input_image.change(
+            fn=apply_image_adjustments,
+            inputs=[input_image, preprocess_invert, gamma_factor],
+            outputs=adjusted_image,
+        )
+
+        gamma_factor.change(
+            fn=apply_image_adjustments,
+            inputs=[input_image, preprocess_invert, gamma_factor],
+            outputs=adjusted_image,
+        )
+
+        preprocess_invert.change(
+            fn=apply_image_adjustments,
+            inputs=[input_image, preprocess_invert, gamma_factor],
+            outputs=adjusted_image,
+        )
 
         # Button to trigger prediction
         submit_button = gr.Button("Submit")
 
-        # Define what happens when the button is clicked
+        # Define what happens when the button is clicked (send adjusted image to predict)
         submit_button.click(
             vsgradio.predict,
-            inputs=input_image,
+            inputs=[adjusted_image, cell_diameter],
             outputs=[output_nucleus, output_membrane],
         )
 
         # Example images and article
         gr.Examples(
-            examples=["examples/a549.png", "examples/hek.png"], inputs=input_image
+            examples=[
+                "examples/a549.png",
+                "examples/hek.png",
+                "examples/ctc_HeLa.png",
+                "examples/livecell_A172.png",
+            ],
+            inputs=input_image,
         )
 
         # Article or footer information
         gr.HTML(
             """
             <div class='article-block'>
-            <p> Model trained primarily on HEK293T, BJ5, and A549 cells. For best results, use quantitative phase images (QPI) or Zernike phase contrast.</p>
-            <p> For training, inference and evaluation of the model refer to the <a href='https://github.com/mehta-lab/VisCy/tree/main/examples/virtual_staining/dlmbl_exercise' target='_blank'>GitHub repository</a>.</p>
+            <p> Model trained primarily on HEK293T, BJ5, and A549 cells. For best results, use quantitative phase images (QPI)</p>
             </div>
             """
         )

requirements.txt

@@ -1,3 +1,4 @@
 viscy<0.3.0
 gradio
-scikit-image
+scikit-image
+cmap