Update app.py

app.py

@@ -118,13 +118,14 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=
     if label_components:
         structure = generate_binary_structure(3, 3)  # 3D, 26-connectivity
         labeled_data, num_features = label(data > 0, structure=structure)
+        labeled_data_resampled, num_features = label(resampled_data > 0, structure=structure)
     else:
         labeled_data = resampled_data
         num_features = None  # Not needed if we're not labeling
 
     # Compute or reuse the center of mass
     if center is None:
-        com = center_of_mass(resampled_data > 0)
+        com = center_of_mass(labeled_data_resampled > 0)
         center = np.round(com).astype(int)
 
     # Define half the slice size
@@ -151,9 +152,9 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=
         return padded_slice
 
     # Extract slices
-    axial_slice = extract_2d_slice(resampled_data, center, axis=2)
-    coronal_slice = extract_2d_slice(resampled_data, center, axis=1)
-    sagittal_slice = extract_2d_slice(resampled_data, center, axis=0)
+    axial_slice = extract_2d_slice(labeled_data_resampled, center, axis=2)
+    coronal_slice = extract_2d_slice(labeled_data_resampled, center, axis=1)
+    sagittal_slice = extract_2d_slice(labeled_data_resampled, center, axis=0)
 
     # Apply rotations
     axial_slice = np.rot90(axial_slice, k=-1)