Update app.py

app.py

@@ -22,16 +22,24 @@ def gradio_demo(model_name, sequence_input, image):
  dataset = "HPA"
  
  
- nucleus_image = image['image']
- protein_image = image['mask']
+ 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
  
- nucleus_image = process_image(nucleus_image, dataset, "nucleus")
- protein_image = process_image(protein_image, dataset, "nucleus")
- protein_image = 1.0*(protein_image > .5)
- print(f'{nucleus_image=}')
- print(f'{protein_image.shape=}')
+ print(f'{protein_image.sum()}')
  
- threshold, heatmap = run_sequence_prediction(
+
+ formatted_predicted_sequence = run_sequence_prediction(
  sequence_input=sequence_input,
  nucleus_image=nucleus_image,
  protein_image=protein_image,
@@ -40,35 +48,16 @@ def gradio_demo(model_name, sequence_input, image):
  device=device,
  )
  
- protein_image = protein_image[0, 0]
- protein_image = protein_image * 1.0
- 
-
- # 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,
- )
+ return formatted_predicted_sequence
 
 
 with gr.Blocks() 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_2560 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_OpenCell_2560 is trained on OpenCell and is good more live-cell predictions on HEK cells."
  )
  with gr.Row():
  model_name = gr.Dropdown(
@@ -88,14 +77,14 @@ with gr.Blocks() as demo:
  )
  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(
  source="upload", 
  tool="sketch",
+ invert_colors=True,
  label="Nucleus Image", 
  line_color="white",
  interactive=True,
@@ -104,12 +93,11 @@ with gr.Blocks() as demo:
  )
 
  with gr.Row():
- gr.Markdown("Image predictions are show below.")
+ gr.Markdown("Sequence predictions are show below.")
 
  with gr.Row().style(equal_height=True):
- predicted_sequence = gr.Textbox(
- label="Predicted Sequence",
- )
+ predicted_sequence = gr.Textbox(label='Predicted Sequence')
+
 
  with gr.Row():
  button = gr.Button("Run Model")
@@ -120,4 +108,4 @@ with gr.Blocks() as demo:
 
  button.click(gradio_demo, inputs, outputs)
 
-demo.launch(share=True)
+demo.launch(enable_queue=True)