Update app.py

app.py

@@ -41,10 +41,10 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     com = center_of_mass(data > 0)
     center = np.round(com).astype(int)
 
-    # Define half the slice size to extract regions around the center of mass
+    # Define half the slice size
     half_size = slice_size // 2
 
-    # Safely extract and crop 2D slices
+    # Safely extract and pad 2D slices
     def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
         slices[axis] = center[axis]  # Fix the axis to extract a single slice
@@ -56,7 +56,15 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
             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
+
+        # Pad the slice to ensure 180x180 dimensions
+        pad_height = slice_size - cropped_slice.shape[0]
+        pad_width = slice_size - cropped_slice.shape[1]
+        padded_slice = np.pad(cropped_slice, 
+                              ((pad_height // 2, pad_height - pad_height // 2), 
+                               (pad_width // 2, pad_width - pad_width // 2)),
+                              mode='constant', constant_values=0)
+        return padded_slice
 
     # Extract slices in axial, coronal, and sagittal planes
     axial_slice = extract_2d_slice(data, center, axis=2)  # Axial (z-axis)
@@ -73,14 +81,14 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     # Create subplots
     fig, axes = plt.subplots(1, 3, figsize=(12, 4))
 
-    # Plot each rotated slice
-    axes[0].imshow(axial_slice, cmap="gray", origin="lower", aspect='auto')
+    # Plot each padded and rotated slice
+    axes[0].imshow(axial_slice, cmap="gray", origin="lower")
     axes[0].axis("off")
 
-    axes[1].imshow(coronal_slice, cmap="gray", origin="lower", aspect='auto')
+    axes[1].imshow(coronal_slice, cmap="gray", origin="lower")
     axes[1].axis("off")
 
-    axes[2].imshow(sagittal_slice, cmap="gray", origin="lower", aspect='auto')
+    axes[2].imshow(sagittal_slice, cmap="gray", origin="lower")
     axes[2].axis("off")
 
     # Save the figure