File size: 5,236 Bytes
5d9796e
1f6e4a8
 
 
5d9796e
1f6e4a8
5d9796e
 
 
 
 
 
 
1f6e4a8
5d9796e
 
 
 
1f6e4a8
5d9796e
 
 
 
 
1f6e4a8
5d9796e
 
 
 
 
 
 
 
 
1f6e4a8
5d9796e
 
1f6e4a8
 
5d9796e
 
90ad424
5d9796e
1f6e4a8
 
5d9796e
1f6e4a8
 
 
 
 
 
5d9796e
 
 
 
1f6e4a8
5d9796e
 
 
1f6e4a8
 
 
 
 
 
5d9796e
 
1f6e4a8
5d9796e
 
 
 
1f6e4a8
5d9796e
 
 
 
 
 
 
 
 
 
 
1f6e4a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import gradio as gr
import torch
from viscy.light.engine import VSUNet
from huggingface_hub import hf_hub_download
from numpy.typing import ArrayLike
import numpy as np
from skimage import exposure


class VSGradio:
    def __init__(self, model_config, model_ckpt_path):
        self.model_config = model_config
        self.model_ckpt_path = model_ckpt_path
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
        self.load_model()

    def load_model(self):
        # Load the model checkpoint and move it to the correct device (GPU or CPU)
        self.model = VSUNet.load_from_checkpoint(
            self.model_ckpt_path,
            architecture="UNeXt2_2D",
            model_config=self.model_config,
        )
        self.model.to(self.device)  # Move the model to the correct device (GPU/CPU)
        self.model.eval()

    def normalize_fov(self, input: ArrayLike):
        "Normalizing the fov with zero mean and unit variance"
        mean = np.mean(input)
        std = np.std(input)
        return (input - mean) / std

    def predict(self, inp):
        # Normalize the input and convert to tensor
        inp = self.normalize_fov(inp)
        inp = torch.from_numpy(np.array(inp).astype(np.float32))

        # Prepare the input dictionary and move input to the correct device (GPU or CPU)
        test_dict = dict(
            index=None,
            source=inp.unsqueeze(0).unsqueeze(0).unsqueeze(0).to(self.device),
        )

        # Run model inference
        with torch.inference_mode():
            self.model.on_predict_start()  # Necessary preprocessing for the model
            pred = (
                self.model.predict_step(test_dict, 0, 0).cpu().numpy()
            )  # Move output back to CPU for post-processing

        # 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


# Load the custom CSS from the file
def load_css(file_path):
    with open(file_path, "r") as file:
        return file.read()


# %%
if __name__ == "__main__":
    # Download the model checkpoint from Hugging Face
    model_ckpt_path = hf_hub_download(
        repo_id="compmicro-czb/VSCyto2D", filename="epoch=399-step=23200.ckpt"
    )

    # Model configuration
    model_config = {
        "in_channels": 1,
        "out_channels": 2,
        "encoder_blocks": [3, 3, 9, 3],
        "dims": [96, 192, 384, 768],
        "decoder_conv_blocks": 2,
        "stem_kernel_size": [1, 2, 2],
        "in_stack_depth": 1,
        "pretraining": False,
    }

    # 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>
            </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")
            with gr.Column():
                output_nucleus = gr.Image(type="numpy", label="VS Nucleus")
                output_membrane = gr.Image(type="numpy", label="VS Membrane")

        # Button to trigger prediction
        submit_button = gr.Button("Submit")

        # Define what happens when the button is clicked
        submit_button.click(
            vsgradio.predict,
            inputs=input_image,
            outputs=[output_nucleus, output_membrane],
        )

        # Example images and article
        gr.Examples(
            examples=["examples/a549.png", "examples/hek.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>
            </div>
            """
        )

    # Launch the Gradio app
    demo.launch()