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 for citations
 # Markdown for citations
 markdown = gr.Markdown("""
 **If you find this tool useful, please consider citing:**
@@ -244,12 +245,26 @@ markdown = gr.Markdown("""
    DOI: [10.1038/s41592-020-01008-z](https://www.nature.com/articles/s41592-020-01008-z)
 """)
 
-# Use Gradio Blocks for layout
+# Use Gradio Blocks for a clean layout
 with gr.Blocks() as demo:
+    # Title and Description
+    gr.Markdown("""
+    # FLAMeS: Multiple Sclerosis Lesion Segmentation
+
+    Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of multiple sclerosis lesions.
+    """)
+    
+    # Layout for Inputs and Outputs
     with gr.Row():
-        interface.render()
-    with gr.Row():
-        markdown.render()
+        with gr.Column(scale=1):  # Input column
+            flair_input = gr.File(label="Upload FLAIR Image (.nii.gz)")
+        with gr.Column(scale=2):  # Output column
+            seg_output = gr.File(label="Download Segmentation Mask")
+            input_img = gr.Image(label="Input: FLAIR image")
+            output_img = gr.Image(label="Output: Lesion Mask")
+    
+    # References
+    gr.Markdown(markdown)
 
 # Debugging GPU environment
 if torch.cuda.is_available():
@@ -258,7 +273,8 @@ else:
     print("No GPU available. Falling back to CPU.")
     os.system("nvidia-smi")  # Check if NVIDIA tools are available
 
-# Launch the app
-if __name__ == "__main__":
-    demo.launch(share=True)
+# Define interface and integrate with Blocks
+def predict_wrapper(file):
+    return run_nnunet_predict(file)
 
+demo.launch(share=True)