Update app.py

app.py

@@ -96,11 +96,14 @@ 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):
+def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=None, label_components=False):
     """
-    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.
+    Extracts slices from a 3D NIfTI image. 
+    If label_components=True, it assigns different labels (colors) to each connected component (26-connectivity)
+    and returns the labeled 3D mask.
+    
+    Returns:
+        labeled_data (np.ndarray): The 3D array (either labeled or original).
     """
     # Load NIfTI image
     img = nib.load(nifti_path)
@@ -110,65 +113,80 @@ def extract_middle_slices(nifti_path, output_image_path, slice_size=180, center=
     # Resample the image to 1 mm isotropic
     resampled_data, _ = resample_to_isotropic(data, affine, target_spacing=1.0)
 
+    # Optionally label connected components
+    if label_components:
+        structure = generate_binary_structure(3, 3)  # 3D, 26-connectivity
+        labeled_data, 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 > 0)
         center = np.round(com).astype(int)
 
     # Define half the slice size
     half_size = slice_size // 2
 
-    # Safely extract and pad 2D slices
+    # Function to extract and pad slices
     def extract_2d_slice(data, center, axis):
         slices = [slice(None)] * 3
-        slices[axis] = center[axis]  # Fix the axis to extract a single slice
+        slices[axis] = center[axis]
         extracted_slice = data[tuple(slices)]
 
-        # Crop the 2D slice around the center in the remaining dimensions
         remaining_axes = [i for i in range(3) if i != axis]
         cropped_slice = extracted_slice[
             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]),
         ]
 
-        # 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), 
+        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(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)
+    # Extract slices
+    axial_slice = extract_2d_slice(labeled_data, center, axis=2)
+    coronal_slice = extract_2d_slice(labeled_data, center, axis=1)
+    sagittal_slice = extract_2d_slice(labeled_data, center, axis=0)
 
-    # Apply rotations to each slice
-    axial_slice = np.rot90(axial_slice, k=-1)  # 90 degrees clockwise
-    coronal_slice = np.rot90(coronal_slice, k=1)  # 90 degrees anticlockwise
-    coronal_slice = np.rot90(coronal_slice, k=2)  # Additional 180 degrees
-    sagittal_slice = np.rot90(sagittal_slice, k=1)  # 90 degrees anticlockwise
-    sagittal_slice = np.rot90(sagittal_slice, k=2)  # Additional 180 degrees
+    # Apply rotations
+    axial_slice = np.rot90(axial_slice, k=-1)
+    coronal_slice = np.rot90(coronal_slice, k=1)
+    coronal_slice = np.rot90(coronal_slice, k=2)
+    sagittal_slice = np.rot90(sagittal_slice, k=1)
+    sagittal_slice = np.rot90(sagittal_slice, k=2)
 
     # Create subplots
     fig, axes = plt.subplots(1, 3, figsize=(12, 4))
 
-    # 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")
-    axes[1].axis("off")
+    # Choose colormap
+    if label_components:
+        cmap = plt.cm.nipy_spectral  # Colorful
+        vmin = 0
+        vmax = num_features
+    else:
+        cmap = "gray"  # Normal
+        vmin = None
+        vmax = None
 
-    axes[2].imshow(sagittal_slice, cmap="gray", origin="lower")
-    axes[2].axis("off")
+    # Plot slices
+    for idx, slice_data in enumerate([axial_slice, coronal_slice, sagittal_slice]):
+        ax = axes[idx]
+        im = ax.imshow(slice_data, cmap=cmap, origin="lower", vmin=vmin, vmax=vmax)
+        ax.axis("off")
 
-    # Save the figure
+    # Save figure
     plt.tight_layout()
     plt.savefig(output_image_path, bbox_inches="tight", pad_inches=0)
     plt.close()
+
+    # Return the labeled mask
+    return labeled_data
     
 # Function to run nnUNet inference
 @spaces.GPU(duration=90)  # Decorate the function to allocate GPU for its execution
@@ -239,8 +257,8 @@ def run_nnunet_predict(nifti_file,hd_bet=False):
             # 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, center=center)
-            extract_middle_slices(new_output_file, output_slice_path, center=center)
+            image = extract_middle_slices(input_path, input_slice_path, center=center)
+            labeled_mask = extract_middle_slices(new_output_file, output_slice_path, center=center, label_components=True)
 
             # Return paths for the Gradio interface
             return new_output_file, input_slice_path, output_slice_path