Update app.py

app.py

@@ -1,119 +1,156 @@
+import os
 import gradio as gr
-from prediction import run_sequence_prediction
+from prediction import run_image_prediction
 import torch
 import torchvision.transforms as T
 from celle.utils import process_image
 from PIL import Image
 from matplotlib import pyplot as plt
+from celle_main import instantiate_from_config
+from omegaconf import OmegaConf
 
 
-def gradio_demo(model_name, sequence_input, image):
-    model = f"CELL-E_2-Image_Prediction/models/{model_name}.ckpt"
-    config = f"CELL-E_2-Image_Prediction/models/{model_name}.yaml"
+class model:
+    def __init__(self):
+        self.model = None
+        self.model_name = None
 
-    if "Finetuned" in model_name:
-        dataset = "OpenCell"
-
-    else:
-        dataset = "HPA"
-        
+    def gradio_demo(self, model_name, sequence_input, nucleus_image, protein_image):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-    nucleus_image = image['image'].convert('L')
-    protein_image = image['mask'].convert('L')
-
-    to_tensor = T.ToTensor()
-    nucleus_tensor = to_tensor(nucleus_image)
-    protein_tensor = to_tensor(protein_image)
-    stacked_images = torch.stack([nucleus_tensor, protein_tensor], dim=0)
-    processed_images = process_image(stacked_images, dataset)
-
-    nucleus_image = processed_images[0].unsqueeze(0)
-    protein_image = processed_images[1].unsqueeze(0)
-    protein_image = protein_image > 0
-    protein_image = 1.0 * protein_image
-    
-    print(f'{protein_image.sum()}')
-    
-
-    formatted_predicted_sequence = run_sequence_prediction(
-        sequence_input=sequence_input,
-        nucleus_image=nucleus_image,
-        protein_image=protein_image,
-        model_ckpt_path=model,
-        model_config_path=config,
-        device=device,
-    )
-    
-    return T.ToPILImage()(protein_image), T.ToPILImage()(nucleus_image), formatted_predicted_sequence
+        if self.model_name != model_name:
+            self.model_name = model_name
+            model_ckpt_path = f"CELL-E_2-Image_Prediction/models/{model_name}.ckpt"
+            model_config_path = f"CELL-E_2-Image_Prediction/models/{model_name}.yaml"
+            
+            
+            # Load model config and set ckpt_path if not provided in config
+            config = OmegaConf.load(model_config_path)
+            if config["model"]["params"]["ckpt_path"] is None:
+                config["model"]["params"]["ckpt_path"] = model_ckpt_path
+
+            # Set condition_model_path and vqgan_model_path to None
+            config["model"]["params"]["condition_model_path"] = None
+            config["model"]["params"]["vqgan_model_path"] = None
+            
+            base_path = os.getcwd()
+
+            os.chdir(os.path.dirname(model_ckpt_path))
+
+            # Instantiate model from config and move to device
+            self.model = instantiate_from_config(config.model).to(device)
+            self.model = torch.compile(self.model,mode='reduce-overhead')
+            
+            os.chdir(base_path)
+
+
+        if "Finetuned" in model_name:
+            dataset = "OpenCell"
+
+        else:
+            dataset = "HPA"
+
+        nucleus_image = process_image(nucleus_image, dataset, "nucleus")
+        if protein_image:
+            protein_image = process_image(protein_image, dataset, "protein")
+            protein_image = protein_image > torch.median(protein_image)
+            protein_image = protein_image[0, 0]
+            protein_image = protein_image * 1.0
+        else:
+            protein_image = torch.ones((256, 256))
+
+        threshold, heatmap = run_image_prediction(
+            sequence_input=sequence_input,
+            nucleus_image=nucleus_image,
+            model=self.model,
+            device=device,
+        )
+
+        # Plot the heatmap
+        plt.imshow(heatmap.cpu(), cmap="rainbow", interpolation="bicubic")
+        plt.axis("off")
+
+        # Save the plot to a temporary file
+        plt.savefig("temp.png", bbox_inches="tight", dpi=256)
+
+        # Open the temporary file as a PIL image
+        heatmap = Image.open("temp.png")
+
+        return (
+            T.ToPILImage()(nucleus_image[0, 0]),
+            T.ToPILImage()(protein_image),
+            T.ToPILImage()(threshold),
+            heatmap,
+        )
 
+base_class = model()
 
 with gr.Blocks(theme='gradio/soft') as demo:
     gr.Markdown("Select the prediction model.")
     gr.Markdown(
-        "- CELL-E_2_HPA_2560 is a good general purpose model for various cell types using ICC-IF."
+        "CELL-E_2_HPA_480 is a good general purpose model for various cell types using ICC-IF."
     )
     gr.Markdown(
-        "- CELL-E_2_OpenCell_2560 is trained on OpenCell and is good more live-cell predictions on HEK cells."
+        "CELL-E_2_HPA_Finetuned_480 is finetuned on OpenCell and is good more live-cell predictions on HEK cells."
     )
     with gr.Row():
         model_name = gr.Dropdown(
-            ["CELL-E_2_HPA_2560", "CELL-E_2_OpenCell_2560"],
-            value="CELL-E_2_HPA_2560",
+            ["CELL-E_2_HPA_480", "CELL-E_2_HPA_Finetuned_480"],
+            value="CELL-E_2_HPA_480",
             label="Model Name",
         )
     with gr.Row():
         gr.Markdown(
-            "Input the desired amino acid sequence. GFP is shown below by default. The sequence must include ```<mask>``` for a prediction to be run."
+            "Input the desired amino acid sequence. GFP is shown below by default."
         )
 
     with gr.Row():
         sequence_input = gr.Textbox(
-            value="M<mask><mask><mask><mask><mask>SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
+            value="MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
             label="Sequence",
         )
     with gr.Row():
         gr.Markdown(
-            "Uploading a nucleus image is necessary. A random crop of 256 x 256 will be applied if larger. We provide default images in [images](https://huggingface.co/spaces/HuangLab/CELL-E_2/tree/main/images). Draw the desired localization on top of the nucelus image."
+            "Uploading a nucleus image is necessary. A random crop of 256 x 256 will be applied if larger. We provide default images in [images](https://huggingface.co/spaces/HuangLab/CELL-E_2/tree/main/images)"
         )
+        gr.Markdown("The protein image is optional and is just used for display.")
 
     with gr.Row().style(equal_height=True):
         nucleus_image = gr.Image(
-            source="upload", 
-            tool="sketch",
-            invert_colors=True,
-            label="Nucleus Image", 
-            interactive=True,
+            type="pil",
+            label="Nucleus Image",
             image_mode="L",
-            type="pil"
         )
 
+        protein_image = gr.Image(type="pil", label="Protein Image (Optional)")
 
-    with gr.Row().style(equal_height=True):
-        nucleus_crop = gr.Image(
-            label="Nucleus Image (Crop)", 
-            image_mode="L",
-            type="pil"
-        )
-
-        mask = gr.Image(
-            label="Threshold Image", 
-            image_mode="L",
-            type="pil"
-        )
     with gr.Row():
-        gr.Markdown("Sequence predictions are show below.")
+        gr.Markdown("Image predictions are show below.")
 
     with gr.Row().style(equal_height=True):
-        predicted_sequence = gr.Textbox(label='Predicted Sequence')
+        nucleus_image_crop = gr.Image(type="pil", label="Nucleus Image", image_mode="L")
+
+        protein_threshold_image = gr.Image(
+            type="pil", label="Protein Threshold Image", image_mode="L"
+        )
 
+        predicted_threshold_image = gr.Image(
+            type="pil", label="Predicted Threshold image", image_mode="L"
+        )
 
+        predicted_heatmap = gr.Image(type="pil", label="Predicted Heatmap")
     with gr.Row():
         button = gr.Button("Run Model")
 
-        inputs = [model_name, sequence_input, nucleus_image]
+        inputs = [model_name, sequence_input, nucleus_image, protein_image]
 
-        outputs = [mask, nucleus_crop, predicted_sequence]
+        outputs = [
+            nucleus_image_crop,
+            protein_threshold_image,
+            predicted_threshold_image,
+            predicted_heatmap,
+        ]
 
-        button.click(gradio_demo, inputs, outputs)
+        button.click(base_class.gradio_demo, inputs, outputs)
 
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True)