Updated app.py

app.py

@@ -29,16 +29,30 @@ def download_model():
 def extract_middle_slice(nifti_path, output_image_path):
     """
     Extracts a middle slice from a 3D NIfTI image and saves it as a PNG file.
+    The figure size is adjusted dynamically based on the slice's aspect ratio.
     """
+    import nibabel as nib
+    import matplotlib.pyplot as plt
+
+    # Load NIfTI image and get the data
     img = nib.load(nifti_path)
     data = img.get_fdata()
-    middle_slice_index = data.shape[2] // 2  # Middle slice along the z-axis
-    plt.figure(figsize=(6, 6))
-    plt.imshow(data[:, :, middle_slice_index], cmap="gray")
+    
+    # Get the middle slice along the z-axis
+    middle_slice_index = data.shape[2] // 2
+    slice_data = data[:, :, middle_slice_index]
+
+    # Calculate aspect ratio
+    height, width = slice_data.shape
+    aspect_ratio = width / height
+
+    # Dynamically adjust figure size based on aspect ratio
+    plt.figure(figsize=(6 * aspect_ratio, 6))  # Height fixed to 6, width scaled by aspect ratio
+    plt.imshow(slice_data, cmap="gray")
     plt.axis("off")
     plt.savefig(output_image_path, bbox_inches="tight", pad_inches=0)
     plt.close()
-    
+
 # Function to run nnUNet inference
 @spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):
@@ -100,9 +114,9 @@ interface = gr.Interface(
     fn=run_nnunet_predict,
     inputs=gr.File(label="Upload FLAIR Image (.nii.gz)"),
     outputs=[
+        gr.File(label="Download Segmentation Mask"),
         gr.Image(label="Input Middle Slice"),
-        gr.Image(label="Output Middle Slice"),
-        gr.File(label="Download Segmentation Mask")
+        gr.Image(label="Output Middle Slice")
     ],
     title="FLAMeS: Multiple Sclerosis Lesion Segmentation",
     description="Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of MS lesions."