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,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
 
-    # Safely extract 2D slices
+    # Safely extract and crop 2D slices
     def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
         slices[axis] = center[axis]
         extracted_slice = data[tuple(slices)]
 
-        # 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)
+        # Crop around the center for the remaining dimensions
+        remaining_axes = [i for i in range(3) if i != axis]
+        for dim in remaining_axes:
+            start = max(center[dim] - half_size, 0)
+            end = min(center[dim] + half_size, extracted_slice.shape[dim])
+            extracted_slice = extracted_slice.take(range(start, end), axis=dim - (dim > axis))
         return extracted_slice
 
     axial_slice = extract_2d_slice(data, center, axis=2)  # Axial (z-axis)