Update app.py

app.py

@@ -1,29 +1,23 @@
-import os
 import gradio as gr
-from prediction import run_image_prediction
+from prediction import run_sequence_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
 
-
 class model:
     def __init__(self):
         self.model = None
         self.model_name = None
 
-    def gradio_demo(self, model_name, sequence_input, nucleus_image, protein_image):
+    def gradio_demo(self, model_name, sequence_input, image):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
         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:
@@ -42,115 +36,110 @@ class model:
             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 = 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()}')
+        
 
-        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(
+        formatted_predicted_sequence = run_sequence_prediction(
             sequence_input=sequence_input,
             nucleus_image=nucleus_image,
-            model=self.model,
+            protein_image=protein_image,
+            model_ckpt_path=self.model,
             device=device,
         )
+        
+        return T.ToPILImage()(protein_image), T.ToPILImage()(nucleus_image), formatted_predicted_sequence
 
-        # 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_480 is a good general purpose model for various cell types using ICC-IF."
+        "- CELL-E_2_HPA_2560 is a good general purpose model for various cell types using ICC-IF."
     )
     gr.Markdown(
-        "CELL-E_2_HPA_Finetuned_480 is finetuned on OpenCell and is good more live-cell predictions on HEK cells."
+        "- CELL-E_2_OpenCell_2560 is trained on OpenCell and is good more live-cell predictions on HEK cells."
     )
     with gr.Row():
         model_name = gr.Dropdown(
-            ["CELL-E_2_HPA_480", "CELL-E_2_HPA_Finetuned_480"],
-            value="CELL-E_2_HPA_480",
+            ["CELL-E_2_HPA_2560", "CELL-E_2_OpenCell_2560"],
+            value="CELL-E_2_HPA_2560",
             label="Model Name",
         )
     with gr.Row():
         gr.Markdown(
-            "Input the desired amino acid sequence. GFP is shown below by default."
+            "Input the desired amino acid sequence. GFP is shown below by default. The sequence must include ```<mask>``` for a prediction to be run."
         )
 
     with gr.Row():
         sequence_input = gr.Textbox(
-            value="MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
+            value="M<mask><mask><mask><mask><mask>SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
             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)"
+            "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."
         )
-        gr.Markdown("The protein image is optional and is just used for display.")
 
     with gr.Row().style(equal_height=True):
         nucleus_image = gr.Image(
-            type="pil",
-            label="Nucleus Image",
+            source="upload", 
+            tool="sketch",
+            invert_colors=True,
+            label="Nucleus Image", 
+            interactive=True,
             image_mode="L",
+            type="pil"
         )
 
-        protein_image = gr.Image(type="pil", label="Protein Image (Optional)")
-
-    with gr.Row():
-        gr.Markdown("Image predictions are show below.")
 
     with gr.Row().style(equal_height=True):
-        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"
+        nucleus_crop = gr.Image(
+            label="Nucleus Image (Crop)", 
+            image_mode="L",
+            type="pil"
         )
 
-        predicted_threshold_image = gr.Image(
-            type="pil", label="Predicted Threshold image", image_mode="L"
+        mask = gr.Image(
+            label="Threshold Image", 
+            image_mode="L",
+            type="pil"
         )
+    with gr.Row():
+        gr.Markdown("Sequence predictions are show below.")
+
+    with gr.Row().style(equal_height=True):
+        predicted_sequence = gr.Textbox(label='Predicted Sequence')
+
 
-        predicted_heatmap = gr.Image(type="pil", label="Predicted Heatmap")
     with gr.Row():
         button = gr.Button("Run Model")
 
-        inputs = [model_name, sequence_input, nucleus_image, protein_image]
+        inputs = [model_name, sequence_input, nucleus_image]
 
-        outputs = [
-            nucleus_image_crop,
-            protein_threshold_image,
-            predicted_threshold_image,
-            predicted_heatmap,
-        ]
+        outputs = [mask, nucleus_crop, predicted_sequence]
 
-        button.click(base_class.gradio_demo, inputs, outputs)
+        button.click(gradio_demo, inputs, outputs)
 
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True)