Update app.py

app.py

@@ -48,7 +48,7 @@ def gradio_demo(model_name, sequence_input, image):
         device=device,
     )
     
-    return formatted_predicted_sequence
+    return T.ToPILImage()(protein_image), formatted_predicted_sequence
 
 
 with gr.Blocks() as demo:
@@ -88,9 +88,16 @@ with gr.Blocks() as demo:
             label="Nucleus Image", 
             interactive=True,
             image_mode="L",
-            type="numpy"
+            type="pil"
         )
 
+        mask = gr.Image(
+            label="Threshold Image", 
+            image_mode="L",
+            type="pil"
+        )
+
+
     with gr.Row():
         gr.Markdown("Sequence predictions are show below.")
 
@@ -103,7 +110,7 @@ with gr.Blocks() as demo:
 
         inputs = [model_name, sequence_input, nucleus_image]
 
-        outputs = [predicted_sequence]
+        outputs = [mask, predicted_sequence]
 
         button.click(gradio_demo, inputs, outputs)