Update app.py

app.py

@@ -33,7 +33,7 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     Extracts slices centered around the center of mass of non-zero voxels in a 3D NIfTI image.
     The slices are taken along axial, coronal, and sagittal planes and saved as a single PNG.
     """
-    # Load NIfTI image and get the data
+      # Load NIfTI image and get the data
     img = nib.load(nifti_path)
     data = img.get_fdata()
 
@@ -44,18 +44,23 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     # Define half the slice size to extract regions around the center of mass
     half_size = slice_size // 2
 
-    # Safely extract slices with boundary checks
-    def safe_slice(data, center, axis, half_size):
+    # Safely extract 2D slices
+    def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
-        slices[axis] = slice(
-            max(center[axis] - half_size, 0), 
-            min(center[axis] + half_size, data.shape[axis])
-        )
-        return data[tuple(slices)]
+        slices[axis] = center[axis]
+        extracted_slice = data[tuple(slices)]
 
-    axial_slice = safe_slice(data, center, axis=2, half_size=half_size)  # Axial (z-axis)
-    coronal_slice = safe_slice(data, center, axis=1, half_size=half_size)  # Coronal (y-axis)
-    sagittal_slice = safe_slice(data, center, axis=0, half_size=half_size)  # Sagittal (x-axis)
+        # Crop around the center for the other two dimensions
+        other_axes = [i for i in range(3) if i != axis]
+        for i in other_axes:
+            start = max(center[i] - half_size, 0)
+            end = min(center[i] + half_size, data.shape[i])
+            extracted_slice = np.take(extracted_slice, range(start, end), axis=i)
+        return extracted_slice
+
+    axial_slice = extract_2d_slice(data, center, axis=2)  # Axial (z-axis)
+    coronal_slice = extract_2d_slice(data, center, axis=1)  # Coronal (y-axis)
+    sagittal_slice = extract_2d_slice(data, center, axis=0)  # Sagittal (x-axis)
 
     # Create subplots
     fig, axes = plt.subplots(1, 3, figsize=(12, 4))
@@ -78,7 +83,6 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     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):