Update app.py

app.py

@@ -27,38 +27,49 @@ 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):
+def extract_middle_slices(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
-    and scaled to be 50% smaller.
+    Extracts middle slices from a 3D NIfTI image in axial, coronal, and sagittal planes
+    and saves them as a single PNG file with subplots.
     """
-    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()
-    
-    # Get the middle slice along the z-axis
-    middle_slice_index = data.shape[2] // 2
-    slice_data = data[:, :, middle_slice_index]
-    
-    # Rotate the slice 90 degrees clockwise
-    slice_data = np.rot90(slice_data, k=-1)
 
-    # Calculate aspect ratio
-    height, width = slice_data.shape
-    aspect_ratio = width / height
+    # Get middle slices along the three planes
+    middle_axial = data[:, :, data.shape[2] // 2]  # Axial (z-axis)
+    middle_coronal = data[:, data.shape[1] // 2, :]  # Coronal (y-axis)
+    middle_sagittal = data[data.shape[0] // 2, :, :]  # Sagittal (x-axis)
+
+    # Rotate slices for proper orientation
+    middle_axial = np.rot90(middle_axial, k=-1)
+    middle_coronal = np.rot90(middle_coronal, k=-1)
+    middle_sagittal = np.rot90(middle_sagittal, k=-1)
+
+    # Create subplots
+    fig, axes = plt.subplots(1, 3, figsize=(12, 4))  # 3 subplots in a row
 
-    # Dynamically adjust figure size based on aspect ratio and scale down by 0.5
-    plt.figure(figsize=(4 * aspect_ratio, 4))  # Height scaled to 3, width scaled proportionally
-    plt.imshow(slice_data, cmap="gray")
-    plt.axis("off")
+    # Plot each slice
+    axes[0].imshow(middle_axial, cmap="gray")
+    axes[0].axis("off")
+    axes[0].set_title("Axial")
+
+    axes[1].imshow(middle_coronal, cmap="gray")
+    axes[1].axis("off")
+    axes[1].set_title("Coronal")
+
+    axes[2].imshow(middle_sagittal, cmap="gray")
+    axes[2].axis("off")
+    axes[2].set_title("Sagittal")
+
+    # Save the figure
+    plt.tight_layout()
     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):