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)