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)