Update app.py

app.py

@@ -32,6 +32,9 @@ def extract_middle_slices(nifti_path, output_image_path):
     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
+    import numpy as np
 
     # Load NIfTI image and get the data
     img = nib.load(nifti_path)
@@ -43,9 +46,16 @@ def extract_middle_slices(nifti_path, output_image_path):
     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)
+    middle_axial = np.rot90(middle_axial, k=-1)  # Rotate 90 degrees clockwise
+    middle_coronal = np.rot90(middle_coronal, k=2)  # Rotate 180 degrees
+    middle_sagittal = np.flipud(np.rot90(middle_sagittal, k=-1))  # Rotate 90 degrees clockwise and flip up-down
+
+    # Pad sagittal slice to match dimensions of other slices
+    max_height = max(middle_axial.shape[0], middle_sagittal.shape[0])
+    max_width = max(middle_axial.shape[1], middle_sagittal.shape[1])
+
+    padded_sagittal = np.zeros((max_height, max_width))
+    padded_sagittal[:middle_sagittal.shape[0], :middle_sagittal.shape[1]] = middle_sagittal
 
     # Create subplots
     fig, axes = plt.subplots(1, 3, figsize=(12, 4))  # 3 subplots in a row
@@ -53,15 +63,12 @@ def extract_middle_slices(nifti_path, output_image_path):
     # 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].imshow(padded_sagittal, cmap="gray")
     axes[2].axis("off")
-    axes[2].set_title("Sagittal")
 
     # Save the figure
     plt.tight_layout()
@@ -69,7 +76,6 @@ def extract_middle_slices(nifti_path, output_image_path):
     plt.close()
 
 
-
 # Function to run nnUNet inference
 @spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):