Update app.py

app.py

@@ -8,7 +8,7 @@ import nibabel as nib
 import matplotlib.pyplot as plt
 import spaces  # Import spaces for GPU decoration
 import numpy as np
-from scipy.ndimage import center_of_mass
+from scipy.ndimage import center_of_mass, zoom
 
 # Define paths
 MODEL_DIR = "./model"  # Local directory to store the downloaded model
@@ -27,18 +27,50 @@ def download_model():
         zip_path = hf_hub_download(repo_id=REPO_ID, filename="Dataset004_WML.zip", cache_dir=MODEL_DIR)
         subprocess.run(["unzip", "-o", zip_path, "-d", MODEL_DIR])
         print("Dataset004_WML downloaded and extracted.")
+        
+def resample_to_isotropic(data, affine, target_spacing=1.0):
+    """
+    Resamples a 3D NIfTI image to isotropic voxel size.
+    
+    Parameters:
+        data (numpy.ndarray): The input 3D image data.
+        affine (numpy.ndarray): The affine transformation matrix.
+        target_spacing (float): Desired isotropic voxel spacing (in mm).
+    
+    Returns:
+        resampled_data (numpy.ndarray): Resampled image data.
+        resampled_affine (numpy.ndarray): Updated affine matrix.
+    """
+    # Extract current voxel dimensions from the affine matrix
+    current_spacing = np.sqrt((affine[:3, :3] ** 2).sum(axis=0))
+
+    # Compute the scaling factors for resampling
+    scaling_factors = current_spacing / target_spacing
+
+    # Resample the data using zoom
+    resampled_data = zoom(data, zoom=scaling_factors, order=1)  # Linear interpolation
 
+    # Update the affine matrix to reflect the new voxel dimensions
+    resampled_affine = affine.copy()
+    resampled_affine[:3, :3] /= scaling_factors[:, np.newaxis]
+
+    return resampled_data, resampled_affine        
+    
 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
     img = nib.load(nifti_path)
     data = img.get_fdata()
+    affine = img.affine
+
+    # Resample the image to 1 mm isotropic
+    resampled_data, _ = resample_to_isotropic(data, affine, target_spacing=1.0)
 
     # Compute the center of mass of non-zero voxels
-    com = center_of_mass(data > 0)
+    com = center_of_mass(resampled_data > 0)
     center = np.round(com).astype(int)
 
     # Define half the slice size
@@ -67,9 +99,9 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
         return padded_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)
+    axial_slice = extract_2d_slice(resampled_data, center, axis=2)  # Axial (z-axis)
+    coronal_slice = extract_2d_slice(resampled_data, center, axis=1)  # Coronal (y-axis)
+    sagittal_slice = extract_2d_slice(resampled_data, center, axis=0)  # Sagittal (x-axis)
 
     # Apply rotations to each slice
     axial_slice = np.rot90(axial_slice, k=-1)  # 90 degrees clockwise
@@ -95,7 +127,7 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     plt.tight_layout()
     plt.savefig(output_image_path, bbox_inches="tight", pad_inches=0)
     plt.close()
-
+    
 # Function to run nnUNet inference
 @spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):