Update app.py

app.py

@@ -27,48 +27,47 @@ def download_model():
         subprocess.run(["unzip", "-o", zip_path, "-d", MODEL_DIR])
         print("Dataset004_WML downloaded and extracted.")
 
-def extract_middle_slices(nifti_path, output_image_path):
+def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     """
-    Extracts middle slices from a 3D NIfTI image in axial, coronal, and sagittal planes
-    and saves them as a single PNG file with subplots.
+    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.
     """
-    import nibabel as nib
-    import matplotlib.pyplot as plt
-    import numpy as np
-
     # Load NIfTI image and get the data
     img = nib.load(nifti_path)
     data = img.get_fdata()
 
-    # Get middle slices along the three planes
-    middle_axial = data[:, :, data.shape[2] // 2]  # Axial (z-axis)
-    middle_coronal = data[:, data.shape[1] // 2, :]  # Coronal (y-axis)
-    middle_sagittal = data[data.shape[0] // 2, :, :]  # Sagittal (x-axis)
+    # Compute the center of mass of non-zero voxels
+    com = center_of_mass(data > 0)
+    center = np.round(com).astype(int)
 
-    # Rotate slices for proper orientation
-    middle_axial = np.rot90(middle_axial, k=-1)  # Rotate 90 degrees clockwise
-    middle_coronal = np.rot90(middle_coronal, k=2)  # Rotate 180 degrees
-    middle_sagittal = np.flipud(np.rot90(middle_sagittal, k=-1))  # Rotate 90 degrees clockwise and flip up-down
+    # Define half the slice size to extract regions around the center of mass
+    half_size = slice_size // 2
 
-    # Pad sagittal slice to match dimensions of other slices
-    max_height = max(middle_axial.shape[0], middle_sagittal.shape[0])
-    max_width = max(middle_axial.shape[1], middle_sagittal.shape[1])
+    # Extract slices around the center of mass
+    def extract_slice(data, center, axis):
+        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)
 
-    padded_sagittal = np.zeros((max_height, max_width))
-    padded_sagittal[:middle_sagittal.shape[0], :middle_sagittal.shape[1]] = middle_sagittal
+    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)
 
     # Create subplots
-    fig, axes = plt.subplots(1, 3, figsize=(12, 4))  # 3 subplots in a row
+    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
 
     # Plot each slice
-    axes[0].imshow(middle_axial, cmap="gray")
+    axes[0].imshow(axial_slice, cmap="gray", origin="lower")
     axes[0].axis("off")
+    axes[0].set_title("Axial")
 
-    axes[1].imshow(middle_coronal, cmap="gray")
+    axes[1].imshow(coronal_slice, cmap="gray", origin="lower")
     axes[1].axis("off")
+    axes[1].set_title("Coronal")
 
-    axes[2].imshow(padded_sagittal, cmap="gray")
+    axes[2].imshow(sagittal_slice, cmap="gray", origin="lower")
     axes[2].axis("off")
+    axes[2].set_title("Sagittal")
 
     # Save the figure
     plt.tight_layout()
@@ -76,6 +75,7 @@ def extract_middle_slices(nifti_path, output_image_path):
     plt.close()
 
 
+
 # Function to run nnUNet inference
 @spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):