Update app.py

app.py

@@ -76,6 +76,28 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180):
     # Define half the slice size
     half_size = slice_size // 2
 
+    def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=None):
+    """
+    Extracts slices from a 3D NIfTI image. If a center is provided, it uses it;
+    otherwise, computes the center of mass of non-zero voxels. Slices are taken 
+    along axial, coronal, and sagittal planes and saved as a single PNG.
+    """
+    # 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 or reuse the center of mass
+    if center is None:
+        com = center_of_mass(resampled_data > 0)
+        center = np.round(com).astype(int)
+
+    # Define half the slice size
+    half_size = slice_size // 2
+
     # Safely extract and pad 2D slices
     def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
@@ -171,11 +193,19 @@ def run_nnunet_predict(nifti_file):
         if os.path.exists(output_file):
             os.rename(output_file, new_output_file)
 
+            # Compute center of mass for the input image
+            img = nib.load(input_path)
+            data = img.get_fdata()
+            affine = img.affine
+            resampled_data, _ = resample_to_isotropic(data, affine, target_spacing=1.0)
+            com = center_of_mass(resampled_data > 0)  # Center of mass
+            center = np.round(com).astype(int)        # Round to integer
+
             # Extract and save 2D slices
             input_slice_path = os.path.join(OUTPUT_DIR, f"{base_filename}_input_slice.png")
             output_slice_path = os.path.join(OUTPUT_DIR, f"{base_filename}_output_slice.png")
-            extract_middle_slices(input_path, input_slice_path)
-            extract_middle_slices(new_output_file, output_slice_path)
+            extract_middle_slices(input_path, input_slice_path, center=center)
+            extract_middle_slices(new_output_file, output_slice_path, center=center)
 
             # Return paths for the Gradio interface
             return new_output_file, input_slice_path, output_slice_path