Updated app.py

app.py

@@ -4,6 +4,8 @@ import os
 import shutil
 from huggingface_hub import hf_hub_download
 import torch
+import nibabel as nib
+import matplotlib.pyplot as plt
 import spaces  # Import spaces for GPU decoration
 
 # Define paths
@@ -24,6 +26,19 @@ def download_model():
         subprocess.run(["unzip", "-o", zip_path, "-d", MODEL_DIR])
         print("Dataset004_WML downloaded and extracted.")
 
+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.
+    """
+    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")
+    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):
@@ -64,18 +79,28 @@ def run_nnunet_predict(nifti_file):
         new_output_file = os.path.join(OUTPUT_DIR, f"{base_filename}_LesionMask.nii.gz")
         if os.path.exists(output_file):
             os.rename(output_file, new_output_file)
-            return new_output_file
+
+            # Extract and save 2D slices
+            input_slice_path = os.path.join(OUTPUT_DIR, f"{base_filename}_input_slice.png")
+            output_slice_path = os.path.join(OUTPUT_DIR, f"{base_filename}_output_slice.png")
+            extract_middle_slice(input_path, input_slice_path)
+            extract_middle_slice(new_output_file, output_slice_path)
+
+            return input_slice_path, output_slice_path, new_output_file
         else:
             return "Error: Output file not found."
     except subprocess.CalledProcessError as e:
         return f"Error: {e}"
 
-
 # Gradio Interface
 interface = gr.Interface(
     fn=run_nnunet_predict,
     inputs=gr.File(label="Upload FLAIR Image (.nii.gz)"),
-    outputs=gr.File(label="Download Segmentation Mask"),
+    outputs=[
+        gr.Image(label="Input Middle Slice"),
+        gr.Image(label="Output Middle Slice"),
+        gr.File(label="Download Segmentation Mask")
+    ],
     title="FLAMeS: Multiple Sclerosis Lesion Segmentation",
     description="Upload a skull-stripped FLAIR image (.nii.gz) to generate a binary segmentation of MS lesions."
 )

requirements.txt

@@ -6,4 +6,5 @@ torchvision
 torchaudio
 nnunetv2
 nibabel
+matplotlib
 numpy