adding library of cells with set parameters

app.py

@@ -46,12 +46,12 @@ class VSGradio:
         new_width = int(width * scale_factor)
         return resize(inp, (new_height, new_width), anti_aliasing=True)
 
-    def predict(self, inp, cell_diameter: float):
+    def predict(self, inp, scaling_factor: 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)
+        inp = apply_rescale_image(inp, scaling_factor)
 
         # Convert the input to a tensor
         inp = torch.from_numpy(np.array(inp).astype(np.float32))
@@ -119,22 +119,28 @@ def apply_image_adjustments(image, invert_image: bool, gamma_factor: float):
     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)
+def apply_rescale_image(image, scaling_factor: float):
+    """Resize the input image according to the scaling factor"""
+    scaling_factor = float(scaling_factor)
     image = resize(
         image,
-        (int(image.shape[0] * scale_factor), int(image.shape[1] * scale_factor)),
+        (int(image.shape[0] * scaling_factor), int(image.shape[1] * scaling_factor)),
         anti_aliasing=True,
     )
     return image
 
 
-# Load the custom CSS from the file
+# Function to clear outputs when a new image is uploaded
+def clear_outputs(image):
+    return (
+        image,
+        None,
+        None,
+    )  # Return None for adjusted_image, output_nucleus, and output_membrane
+
+
 def load_css(file_path):
+    """Load custom CSS"""
     with open(file_path, "r") as file:
         return file.read()
 
@@ -163,7 +169,14 @@ if __name__ == "__main__":
     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>"
+            """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                <a href="https://www.czbiohub.org/sf/" target="_blank">
+                <img src="https://huggingface.co/spaces/compmicro-czb/VirtualStaining/resolve/main/misc/czb_mark.png" style="width: 100px; height: auto; margin-right: 10px;">
+                </a>
+                <div class='title-block'>Image Translation (Virtual Staining) of cellular landmark organelles</div>
+            </div>
+            """
         )
         gr.HTML(
             """
@@ -171,9 +184,11 @@ if __name__ == "__main__":
                 <p><b>Model:</b> VSCyto2D</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><b>Note:</b> The model works well with QPI, and sometimes generalizes to phase contrast and DIC.<br>
+                It was trained primarily on HEK293T, BJ5, and A549 cells imaged at 20x. <br>
+                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>
+                <p> For training your own model and analyzing large amounts of data, use our <a href='https://github.com/mehta-lab/VisCy/tree/main/examples/virtual_staining/dlmbl_exercise' target='_blank'>GitHub repository</a>.</p>
             </div>
             """
         )
@@ -182,7 +197,10 @@ if __name__ == "__main__":
         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)"
+                type="numpy",
+                image_mode="L",
+                label="Adjusted Image (Preview)",
+                interactive=False,
             )
 
             with gr.Column():
@@ -201,20 +219,21 @@ if __name__ == "__main__":
 
         # Slider for gamma adjustment
         gamma_factor = gr.Slider(
-            label="Adjust Gamma", minimum=0.1, maximum=5.0, value=1.0, step=0.1
+            label="Adjust Gamma", minimum=0.01, 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",
+        scaling_factor = gr.Textbox(
+            label="Rescaling image factor",
+            value="1.0",
+            placeholder="Rescaling factor for the input image",
         )
 
         # Checkbox for merging predictions
-        merge_checkbox = gr.Checkbox(label="Merge Predictions into one image", value=False)
+        merge_checkbox = gr.Checkbox(
+            label="Merge Predictions into one image", value=True
+        )
 
-        # Update the adjusted image based on all the transformations
         input_image.change(
             fn=apply_image_adjustments,
             inputs=[input_image, preprocess_invert, gamma_factor],
@@ -226,6 +245,15 @@ if __name__ == "__main__":
             inputs=[input_image, preprocess_invert, gamma_factor],
             outputs=adjusted_image,
         )
+        cell_name = gr.Textbox(
+            label="Cell Name", placeholder="Cell Type", visible=False
+        )
+        imaging_modality = gr.Textbox(
+            label="Imaging Modality", placeholder="Imaging Modality", visible=False
+        )
+        references = gr.Textbox(
+            label="References", placeholder="References", visible=False
+        )
 
         preprocess_invert.change(
             fn=apply_image_adjustments,
@@ -237,27 +265,64 @@ if __name__ == "__main__":
         submit_button = gr.Button("Submit")
 
         # Function to handle prediction and merging if needed
-        def submit_and_merge(inp, cell_diameter, merge):
-            nucleus, membrane = vsgradio.predict(inp, cell_diameter)
+        def submit_and_merge(inp, scaling_factor, merge):
+            nucleus, membrane = vsgradio.predict(inp, scaling_factor)
             if merge:
                 merged = merge_images(nucleus, membrane)
-                return merged, gr.update(visible=True), nucleus, gr.update(visible=False), membrane,gr.update(visible=False)
+                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)
+                return (
+                    None,
+                    gr.update(visible=False),
+                    nucleus,
+                    gr.update(visible=True),
+                    membrane,
+                    gr.update(visible=True),
+                )
 
         submit_button.click(
             fn=submit_and_merge,
-            inputs=[adjusted_image, cell_diameter, merge_checkbox],
-            outputs=[merged_image, merged_image, output_nucleus, output_nucleus,output_membrane, output_membrane],
+            inputs=[adjusted_image, scaling_factor, merge_checkbox],
+            outputs=[
+                merged_image,
+                merged_image,
+                output_nucleus,
+                output_nucleus,
+                output_membrane,
+                output_membrane,
+            ],
+        )
+        # Clear everything when the input image changes
+        input_image.change(
+            fn=clear_outputs,
+            inputs=input_image,
+            outputs=[adjusted_image, output_nucleus, 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)
+                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)
+                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(
@@ -267,21 +332,61 @@ if __name__ == "__main__":
         )
 
         # Example images and article
-        gr.Examples(
+        examples_component = gr.Examples(
             examples=[
-                "examples/a549.png",
-                "examples/hek.png",
-                "examples/ctc_HeLa.png",
-                "examples/livecell_A172.png",
+                ["examples/a549.png", "A549", "QPI", 1.0, False, "1.0", "1"],
+                ["examples/hek.png", "HEK293T", "QPI", 1.0, False, "1.0", "1"],
+                ["examples/HEK_PhC.png", "HEK293T", "PhC", 1.2, True, "1.0", "1"],
+                ["examples/livecell_A172.png", "A172", "PhC", 1.0, True, "1.0", "2"],
+                ["examples/ctc_HeLa.png", "HeLa", "DIC", 0.7, False, "0.7", "3"],
+                [
+                    "examples/ctc_glioblastoma_astrocytoma_U373.png",
+                    "Glioblastoma",
+                    "PhC",
+                    1.0,
+                    True,
+                    "2.0",
+                    "3",
+                ],
+                ["examples/U2OS_BF.png", "U2OS", "Brightfield", 1.0, False, "0.3", "4"],
+                ["examples/U2OS_QPI.png", "U2OS", "QPI", 1.0, False, "0.3", "4"],
+                [
+                    "examples/neuromast2.png",
+                    "Zebrafish neuromast",
+                    "QPI",
+                    0.6,
+                    False,
+                    "1.2",
+                    "1",
+                ],
+                [
+                    "examples/mousekidney.png",
+                    "Mouse Kidney",
+                    "QPI",
+                    0.8,
+                    False,
+                    "0.6",
+                    "1",
+                ],
+            ],
+            inputs=[
+                input_image,
+                cell_name,
+                imaging_modality,
+                gamma_factor,
+                preprocess_invert,
+                scaling_factor,
+                references,
             ],
-            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)</p>
+            <li>1. <a href='https://www.biorxiv.org/content/10.1101/2024.05.31.596901' target='_blank'>Liu et al., Robust virtual staining of landmark organelles</a></li>
+            <li>2. <a href='https://sartorius-research.github.io/LIVECell/' target='_blank'>Edlund et. al. LIVECEll-A large-scale dataset for label-free live cell segmentation</a></li>
+            <li>3. <a href='https://celltrackingchallenge.net/' target='_blank'>Maska et. al.,The cell tracking challenge: 10 years of objective benchmarking </a></li>
+            <li>4. <a href='https://elifesciences.org/articles/55502' target='_blank'>Guo et. al., Revealing architectural order with quantitative label-free imaging and deep learning</a></li>
             </div>
             """
         )