Update app.py

app.py

@@ -48,7 +48,7 @@ def gradio_demo(model_name, sequence_input, image):
         device=device,
     )
     
-    return T.ToPILImage()(protein_image), formatted_predicted_sequence
+    return T.ToPILImage()(protein_image), T.ToPILImage()(nucleus_image), formatted_predicted_sequence
 
 
 with gr.Blocks() as demo:
@@ -91,13 +91,19 @@ with gr.Blocks() as demo:
             type="pil"
         )
 
+
+    with gr.Row().style(equal_height=True):
+        nucleus_crop = gr.Image(
+            label="Nucleus Image (Crop)", 
+            image_mode="L",
+            type="pil"
+        )
+
         mask = gr.Image(
             label="Threshold Image", 
             image_mode="L",
             type="pil"
         )
-
-
     with gr.Row():
         gr.Markdown("Sequence predictions are show below.")
 
@@ -110,7 +116,7 @@ with gr.Blocks() as demo:
 
         inputs = [model_name, sequence_input, nucleus_image]
 
-        outputs = [mask, predicted_sequence]
+        outputs = [mask, nucleus_crop, predicted_sequence]
 
         button.click(gradio_demo, inputs, outputs)