improved contour extraction for thin CC and surface coordinates

Author: ClePol

CorpusCallosum/fastsurfer_cc.py

@@ -100,8 +100,8 @@ def options_parse() -> argparse.Namespace:
     parser.add_argument(
         "--contour_smoothing",
         type=float,
-        default=1.0,
-        help="Gaussian sigma for smoothing during contour detection. Default is 1.0, higher values mean a smoother"
+        default=5,
+        help="Window size for smoothing during contour detection. Default is 5, higher values mean a smoother"
         "outline, at the cost of precision.",
     )
     parser.add_argument(
@@ -390,7 +390,7 @@ def segment_cc(midslices, ac_coords, pc_coords, aseg_nib, model_segmentation, sl
         or np.any(cc_volume_mask[:, :, 0])
         or np.any(cc_volume_mask[:, :, -1])
     ):
-        print("Warning: CC volume mask touches the edge of the segmentation field-of-view, CC might be truncated")
+        print("Warning: CC voume mask touches the edge of the segmentation field-of-view, CC might be truncated")
 
     # get voxels that were removed during cleaning
     removed_voxels = pre_clean_segmentation != segmentation
@@ -411,7 +411,7 @@ def main(
     num_thickness_points: int = 100,
     subdivisions: list[float] | None = None,
     subdivision_method: str = "shape",
-    contour_smoothing: float = 1.0,
+    contour_smoothing: float = 5,
     save_template: str | Path | None = None,
     cpu: bool = False,
     # output paths
@@ -606,7 +606,6 @@ def main(
         midslices, ac_coords, pc_coords, aseg_nib, model_segmentation, slices_to_analyze
     )
 
-
     # calculate affine for segmentation volume
     orig_to_seg = np.eye(4)
     orig_to_seg[0, 3] = -FSAVERAGE_MIDDLE + slices_to_analyze // 2

CorpusCallosum/segmentation/segmentation_postprocessing.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 from scipy import integrate, ndimage
+from scipy.spatial.distance import cdist
 from skimage.measure import label
 
 import FastSurferCNN.utils.logging as logging
@@ -22,6 +23,205 @@
 logger = logging.get_logger(__name__)
 
 
+def find_component_boundaries(labels_arr: np.ndarray, component_id: int) -> np.ndarray:
+    """Find boundary voxels of a connected component.
+    
+    Args:
+        labels_arr (np.ndarray): Labeled array from connected components analysis
+        component_id (int): ID of the component to find boundaries for
+        
+    Returns:
+        np.ndarray: Array of boundary coordinates (N, 3)
+    """
+    component_mask = labels_arr == component_id
+    
+    # Create a structuring element for 6-connectivity (face neighbors only)
+    struct = ndimage.generate_binary_structure(3, 1)
+    
+    # Erode the component to find internal voxels
+    eroded = ndimage.binary_erosion(component_mask, structure=struct)
+    
+    # Boundary is the difference between original and eroded
+    boundary = component_mask & ~eroded
+    
+    return np.array(np.where(boundary)).T
+
+
+def find_minimal_connection_path(boundary1: np.ndarray, boundary2: np.ndarray, 
+                                max_distance: float = 3.0) -> tuple[np.ndarray, np.ndarray] | None:
+    """Find the minimal connection path between two component boundaries.
+    
+    Args:
+        boundary1 (np.ndarray): Boundary coordinates of first component (N1, 3)
+        boundary2 (np.ndarray): Boundary coordinates of second component (N2, 3)
+        max_distance (float): Maximum distance to consider for connection
+        
+    Returns:
+        tuple | None: (point1, point2) coordinates of closest points if within max_distance, None otherwise
+    """
+    if len(boundary1) == 0 or len(boundary2) == 0:
+        return None
+    
+    # Calculate pairwise distances between all boundary points
+    distances = cdist(boundary1, boundary2, metric='euclidean')
+    
+    # Find the minimum distance and corresponding points
+    min_idx = np.unravel_index(np.argmin(distances), distances.shape)
+    min_distance = distances[min_idx]
+    
+    if min_distance <= max_distance:
+        point1 = boundary1[min_idx[0]]
+        point2 = boundary2[min_idx[1]]
+        return point1, point2
+    
+    return None
+
+
+def create_connection_line(point1: np.ndarray, point2: np.ndarray) -> list[tuple[int, int, int]]:
+    """Create a line of voxels connecting two points using simplified 3D line algorithm.
+    
+    Args:
+        point1 (np.ndarray): Starting point coordinates (3,)
+        point2 (np.ndarray): Ending point coordinates (3,)
+        
+    Returns:
+        list: List of (x, y, z) coordinates forming the connection line
+    """
+    x1, y1, z1 = map(int, point1)
+    x2, y2, z2 = map(int, point2)
+    
+    line_points = []
+    
+    # Calculate the number of steps needed
+    dx = abs(x2 - x1)
+    dy = abs(y2 - y1)
+    dz = abs(z2 - z1)
+    
+    steps = max(dx, dy, dz)
+    
+    if steps == 0:
+        return [(x1, y1, z1)]
+    
+    # Calculate increments for each dimension
+    x_inc = (x2 - x1) / steps
+    y_inc = (y2 - y1) / steps
+    z_inc = (z2 - z1) / steps
+    
+    # Generate points along the line
+    for i in range(steps + 1):
+        x = int(round(x1 + i * x_inc))
+        y = int(round(y1 + i * y_inc))
+        z = int(round(z1 + i * z_inc))
+        line_points.append((x, y, z))
+    
+    return line_points
+
+
+def connect_nearby_components(seg_arr: np.ndarray, max_connection_distance: float = 3.0) -> np.ndarray:
+    """Connect nearby disconnected components that should be connected.
+    
+    This function identifies disconnected components in the segmentation and creates
+    minimal connections between components that are close to each other.
+    
+    Args:
+        seg_arr (np.ndarray): Input binary segmentation array
+        max_connection_distance (float): Maximum distance to connect components
+        
+    Returns:
+        np.ndarray: Segmentation array with minimal connections added
+    """
+
+    # Create a copy to modify
+    connected_seg = seg_arr.copy()
+    
+    # Find connected components without dilation first
+    labels_cc = label(seg_arr, connectivity=3, background=0)
+    
+    # Get component sizes (excluding background)
+    bincount = np.bincount(labels_cc.flat)
+    component_ids = np.where(bincount > 0)[0][1:]  # Exclude background (0)
+    
+    if len(component_ids) <= 1:
+        return connected_seg  # Only one component, no connections needed
+    
+    # Sort components by size (largest first)
+    component_sizes = [(comp_id, bincount[comp_id]) for comp_id in component_ids]
+    component_sizes.sort(key=lambda x: x[1], reverse=True)
+    
+    # Use the largest component as the reference
+    main_component_id = component_sizes[0][0]
+
+
+    
+    logger.info(f"Found {len(component_ids)} disconnected components. "
+                f"Attempting to connect smaller components to main component (size: {component_sizes[0][1]})")
+    
+    connections_made = 0
+    
+    # Try to connect each smaller component to the main component
+    for comp_id, comp_size in component_sizes[1:]:
+        if comp_size < 5:  # Skip very small components (likely noise)
+            logger.debug(f"Skipping tiny component {comp_id} with size {comp_size}")
+            continue
+            
+        # Find boundaries of both components
+        main_boundary = find_component_boundaries(labels_cc, main_component_id)
+        comp_boundary = find_component_boundaries(labels_cc, comp_id)
+        
+        # Find minimal connection path
+        connection = find_minimal_connection_path(main_boundary, comp_boundary, max_connection_distance)
+        
+        if connection is not None:
+            point1, point2 = connection
+            distance = np.linalg.norm(point2 - point1)
+            
+            logger.debug(f"Connecting component {comp_id} (size: {comp_size}) to main component. "
+                        f"Distance: {distance:.2f} voxels")
+            
+            # Create connection line
+            connection_line = create_connection_line(point1, point2)
+            
+            # Add connection voxels to the segmentation
+            # Use the same label as the original segmentation at the connection points
+            connection_label = seg_arr[point1[0], point1[1], point1[2]] if \
+                seg_arr[point1[0], point1[1], point1[2]] != 0 else \
+                    seg_arr[point2[0], point2[1], point2[2]]
+            
+            for x, y, z in connection_line:
+                if (0 <= x < connected_seg.shape[0] and 
+                    0 <= y < connected_seg.shape[1] and 
+                    0 <= z < connected_seg.shape[2]):
+                    if connected_seg[x, y, z] == 0:  # Only fill empty voxels
+                        connected_seg[x, y, z] = connection_label
+            
+            connections_made += 1
+        else:
+            logger.debug(f"Component {comp_id} (size: {comp_size}) too far from main component")
+    
+    logger.info(f"Created {connections_made} minimal connections between components")
+
+
+    # Plot components for visualization
+    # import matplotlib.pyplot as plt
+    # n_components = len(component_sizes)
+    # fig, axes = plt.subplots(1, n_components + 1, figsize=(5*(n_components + 1), 5))
+    # if n_components == 1:
+    #     axes = [axes]
+    # # Plot each component in a different color
+    # for i, (comp_id, comp_size) in enumerate(component_sizes):
+    #     component_mask = labels_cc == comp_id
+    #     axes[i].imshow(component_mask[component_mask.shape[0]//2], cmap='gray')
+    #     axes[i].set_title(f'Component {comp_id}\nSize: {comp_size}')
+    #     axes[i].axis('off')
+    
+    # # Plot the connected segmentation
+    # axes[-1].imshow(connected_seg[connected_seg.shape[0]//2], cmap='gray')
+    # axes[-1].set_title('Connected Segmentation')
+    # axes[-1].axis('off')
+    # plt.tight_layout()
+    # plt.show()
+    
+    return connected_seg
 
 
 def get_cc_volume_voxel(desired_width_mm: int, cc_mask: np.ndarray, voxel_size: tuple[float, float, float]) -> float:
@@ -142,72 +342,94 @@ def get_cc_volume_contour(desired_width_mm: int, cc_contours: list[np.ndarray],
     return integrate.simpson(areas, x=measurement_points)
 
 
-def get_largest_cc(seg_arr: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    """Get largest connected component from a binary segmentation array.
+def get_largest_cc(seg_arr: np.ndarray, max_connection_distance: float = 3.0) -> tuple[np.ndarray, np.ndarray]:
+    """Get largest connected component from a binary segmentation array with minimal connections.
     
-    This function takes a binary segmentation array, dilates it, finds connected components,
-    and returns the largest component (excluding background) along with its mask.
+    This function takes a binary segmentation array, attempts to connect nearby disconnected
+    components that should be connected, then finds the largest connected component.
+    It first tries to establish minimal connections between close components before
+    falling back to dilation if no connections are made.
     
     Args:
         seg_arr (np.ndarray): Input binary segmentation array
+        max_connection_distance (float): Maximum distance to connect components (default: 3.0)
         
     Returns:
         tuple: A tuple containing:
             - clean_seg (np.ndarray): Segmentation array with only the largest connected component
             - largest_cc (np.ndarray): Binary mask of the largest connected component
     """
-    # generate dilatation structure
-    struct1 = ndimage.generate_binary_structure(3, 3)
-    # Dilate prediction
-    mask = ndimage.binary_dilation(seg_arr, structure=struct1, iterations=1, ).astype(np.uint8)
-    # Get connected components
+    # First attempt: try to connect nearby components with minimal connections
+    connected_seg = connect_nearby_components(seg_arr, max_connection_distance)
+    
+    # Check if connections were successful by comparing connectivity
+    original_labels = label(seg_arr, connectivity=3, background=0)
+    connected_labels = label(connected_seg, connectivity=3, background=0)
+    
+    original_components = len(np.unique(original_labels)) - 1  # Exclude background
+    connected_components = len(np.unique(connected_labels)) - 1  # Exclude background
+    
+    if connected_components < original_components:
+        logger.info(f"Successfully reduced components from {original_components} to {connected_components} "
+                     "using minimal connections")
+    mask = connected_seg
+    # else:
+    #     logger.info("No connections made, falling back to dilation approach")
+    #     # Fallback: use the original dilation approach
+    #     struct1 = ndimage.generate_binary_structure(3, 3)
+    #     mask = ndimage.binary_dilation(seg_arr, structure=struct1, iterations=1).astype(np.uint8)
+    
+    # Get connected components from the processed mask
     labels_cc = label(mask, connectivity=3, background=0)
-    # Get componnets count
+    
+    # Get component counts
     bincount = np.bincount(labels_cc.flat)
-    # Get background label, assumption that background is the biggest connected component
+    
+    # Get background label (assumed to be the largest component)
     background = np.argmax(bincount)
     bincount[background] = -1
+    
     # Get largest connected component
     largest_cc = labels_cc == np.argmax(bincount)
-    # Apply mask
-    clean_seg = seg_arr * largest_cc
 
-    return clean_seg,largest_cc
+    return largest_cc
 
-def clean_cc_segmentation(seg_arr: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+def clean_cc_segmentation(seg_arr: np.ndarray, max_connection_distance: float = 3.0) -> tuple[np.ndarray, np.ndarray]:
     """Clean corpus callosum segmentation by removing non-connected components.
     
     This function processes a segmentation array to clean up the corpus callosum (CC)
     by removing non-connected components. It first isolates the CC (label 192),
-    removes non-connected components, then adds the fornix (label 250), and
-    finally removes non-connected components from the combined CC and fornix.
+    attempts to connect nearby disconnected components, then adds the fornix (label 250), 
+    and finally removes non-connected components from the combined CC and fornix.
     
     Args:
         seg_arr (np.ndarray): Input segmentation array with CC (192) and fornix (250) labels
+        max_connection_distance (float): Maximum distance to connect components (default: 3.0)
         
     Returns:
         tuple: A tuple containing:
             - clean_seg (np.ndarray): Cleaned segmentation array with only the largest 
               connected component of CC and fornix
             - mask (np.ndarray): Binary mask of the largest connected component
     """
-    #Remove non connected components from the CC alone
-    clean_seg = np.zeros_like(seg_arr)
-    clean_seg[seg_arr == CC_LABEL] = CC_LABEL
-    clean_seg,_ = get_largest_cc(clean_seg)
+    # Remove non connected components from the CC alone, with minimal connections
+    cc_seg = np.zeros_like(seg_arr)
+    cc_seg[seg_arr == CC_LABEL] = CC_LABEL
 
-    #Add fornix to the CC labels
-    clean_seg[seg_arr == FORNIX_LABEL] = FORNIX_LABEL
+    cc_label_cleaned = np.zeros_like(cc_seg)
+    for i in range(cc_seg.shape[0]):
+        cc_label_cleaned[i] = get_largest_cc(cc_seg[None,i], max_connection_distance)
+        # import matplotlib.pyplot as plt
+        # fig, ax = plt.subplots(1,3)
+        # ax[0].imshow(cc_seg[i])
+        # ax[1].imshow(mask[i])
+        # ax[2].imshow(cc_seg[i] - mask[i]*CC_LABEL) # difference between pre and post clean
+        # plt.show()
 
-    #Remove non connected components from CC & Fornix
-    clean_seg, mask = get_largest_cc(clean_seg)
 
-    unique_labels = np.unique(clean_seg)
+    # Add fornix to the CC labels
+    clean_seg = np.zeros_like(seg_arr)
+    clean_seg[cc_label_cleaned > 0] = CC_LABEL
+    clean_seg[seg_arr == FORNIX_LABEL] = FORNIX_LABEL
 
-    if 250 not in unique_labels:
-        clean_seg[seg_arr == 250] = 250
-        mask[seg_arr == 250] = True
-    if 192 not in unique_labels:
-        clean_seg[seg_arr == 192] = 192
-        mask[seg_arr == 192] = True
-    return clean_seg, mask
+    return clean_seg, cc_label_cleaned > 0