Update app.py

app.py

@@ -44,15 +44,18 @@ 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
 
-    # Extract slices around the center of mass
-    def extract_slice(data, center, axis):
+    # Safely extract slices with boundary checks
+    def safe_slice(data, center, axis, half_size):
         slices = [slice(None)] * 3
-        slices[axis] = slice(center[axis] - half_size, center[axis] + half_size)
-        return np.take(data, range(center[axis] - half_size, center[axis] + half_size), axis=axis, mode='constant', cval=0)
+        slices[axis] = slice(
+            max(center[axis] - half_size, 0), 
+            min(center[axis] + half_size, data.shape[axis])
+        )
+        return data[tuple(slices)]
 
-    axial_slice = extract_slice(data, center, axis=2)  # Axial (z-axis)
-    coronal_slice = extract_slice(data, center, axis=1)  # Coronal (y-axis)
-    sagittal_slice = extract_slice(data, center, axis=0)  # Sagittal (x-axis)
+    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)
 
     # Create subplots
     fig, axes = plt.subplots(1, 3, figsize=(12, 4))
@@ -76,7 +79,6 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     plt.close()
 
 
-
 # Function to run nnUNet inference
 @spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):