Update app.py

app.py

@@ -47,17 +47,18 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     # Safely extract and crop 2D slices
     def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
-        slices[axis] = center[axis]
+        slices[axis] = center[axis]  # Fix the axis to extract a single slice
         extracted_slice = data[tuple(slices)]
 
-        # Crop around the center for the remaining dimensions
+        # Crop the 2D slice around the center in 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
+        cropped_slice = extracted_slice[
+            max(center[remaining_axes[0]] - half_size, 0):min(center[remaining_axes[0]] + half_size, extracted_slice.shape[0]),
+            max(center[remaining_axes[1]] - half_size, 0):min(center[remaining_axes[1]] + half_size, extracted_slice.shape[1]),
+        ]
+        return cropped_slice
 
+    # Extract slices in axial, coronal, and sagittal planes
     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)