Update app.py

app.py

@@ -229,6 +229,7 @@ interface = gr.Interface(
     description="Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of multiple sclerosis lesions.",
 )
 
+# Markdown content as a string
 # Markdown content as a string
 markdown_content = """
 **If you find this tool useful, please consider citing:**
@@ -257,6 +258,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=1):  # Input column
             flair_input = gr.File(label="Upload FLAIR Image (.nii.gz)")
+            submit_button = gr.Button("Submit")
         with gr.Column(scale=2):  # Output column
             seg_output = gr.File(label="Download Segmentation Mask")
             input_img = gr.Image(label="Input: FLAIR image")
@@ -265,6 +267,13 @@ with gr.Blocks() as demo:
     # References
     gr.Markdown(markdown_content)
 
+    # Link submit button to the function
+    submit_button.click(
+        fn=run_nnunet_predict, 
+        inputs=[flair_input], 
+        outputs=[seg_output, input_img, output_img]
+    )
+
 # Debugging GPU environment
 if torch.cuda.is_available():
     print(f"GPU is available: {torch.cuda.get_device_name(0)}")
@@ -272,8 +281,6 @@ else:
     print("No GPU available. Falling back to CPU.")
     os.system("nvidia-smi")  # Check if NVIDIA tools are available
 
-# Define interface and integrate with Blocks
-def predict_wrapper(file):
-    return run_nnunet_predict(file)
-
-demo.launch(share=True)
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(share=True)