add binary segmentation decoding and postprocessing

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -628,8 +628,8 @@ class AffineConfig:
     enabled: Optional[bool] = None
     prob: float = 0.5
     rotate_range: Tuple[float, float, float] = (0.2, 0.2, 0.2)  # Rotation range in radians (~11°)
-    scale_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)   # Scaling range (±10%)
-    shear_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)   # Shearing range (±10%)
+    scale_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)  # Scaling range (±10%)
+    shear_range: Tuple[float, float, float] = (0.1, 0.1, 0.1)  # Shearing range (±10%)
 
 
 @dataclass
@@ -913,16 +913,64 @@ class DecodeModeConfig:
     )  # Keyword arguments for the decode function
 
 
+@dataclass
+class BinaryPostprocessingConfig:
+    """Binary segmentation postprocessing configuration.
+
+    Applies morphological operations and connected components filtering to binary predictions.
+    Pipeline order:
+        1. Threshold predictions to binary mask (handled by decoding)
+        2. Apply morphological opening (erosion + dilation)
+        3. Extract connected components
+        4. Keep top-k largest components
+    """
+
+    enabled: bool = False  # Enable binary postprocessing pipeline
+    median_filter_size: Optional[Tuple[int, ...]] = (
+        None  # Median filter kernel size (e.g., (3, 3) for 2D)
+    )
+    opening_iterations: int = 0  # Number of morphological opening iterations
+    closing_iterations: int = 0  # Number of morphological closing iterations
+    connected_components: Optional["ConnectedComponentsConfig"] = None  # CC filtering config
+
+
+@dataclass
+class ConnectedComponentsConfig:
+    """Connected components filtering configuration."""
+
+    enabled: bool = True  # Enable connected components filtering
+    top_k: Optional[int] = None  # Keep only top-k largest components (None = keep all)
+    min_size: int = 0  # Minimum component size in voxels
+    connectivity: int = 1  # Connectivity for CC (1=face, 2=face+edge, 3=face+edge+corner)
+
+
 @dataclass
 class PostprocessingConfig:
-    """Postprocessing configuration."""
+    """Postprocessing configuration for inference output.
+
+    Controls how predictions are transformed before saving:
+    - Thresholding: Binarize predictions using threshold_range
+    - Scaling: Multiply intensity values (e.g., 255 for uint8)
+    - Dtype conversion: Convert to target data type with proper clamping
+    - Transpose: Reorder axes (e.g., [2,1,0] for zyx->xyz)
+    """
 
-    output_scale: Optional[float] = (
+    # Thresholding configuration
+    binary: Optional[Dict[str, Any]] = field(
+        default_factory=dict
+    )  # Binary thresholding config (e.g., {'threshold_range': [0.5, 1.0]})
+
+    # Intensity scaling
+    intensity_scale: Optional[float] = (
         None  # Scale predictions for saving (e.g., 255.0 for uint8). None = no scaling
     )
-    output_dtype: Optional[str] = (
+
+    # Data type conversion
+    intensity_dtype: Optional[str] = (
         None  # Output data type: 'uint8', 'uint16', 'float32'. None = no conversion (keep as-is)
     )
+
+    # Axis permutation
     output_transpose: List[int] = field(
         default_factory=list
     )  # Axis permutation for output (e.g., [2,1,0] for zyx->xyz)
@@ -1072,6 +1120,8 @@ def configure_instance_segmentation(cfg: Config, boundary_thickness: int = 5) ->
     "SlidingWindowConfig",
     "TestTimeAugmentationConfig",
     "PostprocessingConfig",
+    "BinaryPostprocessingConfig",
+    "ConnectedComponentsConfig",
     "EvaluationConfig",
     "DecodeModeConfig",
     "DecodeBinaryContourDistanceWatershedConfig",

connectomics/data/process/__init__.py

@@ -1,6 +1,10 @@
 # Core processing functions
 from .bbox_processor import *  # New: unified bbox processing framework
 
+# Utility functions used by decoding
+from .misc import get_seg_type
+from .bbox import bbox_ND, crop_ND, replace_ND
+
 # MONAI-native transforms and composition
 
 # Pipeline builder (primary entry point for label transforms)

connectomics/decoding/__init__.py

@@ -19,6 +19,7 @@
 """
 
 from .segmentation import (
+    decode_binary_thresholding,
     decode_binary_cc,
     decode_binary_watershed,
     decode_binary_contour_cc,
@@ -58,6 +59,7 @@
 
 __all__ = [
     # Segmentation decoding
+    'decode_binary_thresholding',
     'decode_binary_cc',
     'decode_binary_watershed',
     'decode_binary_contour_cc',

connectomics/decoding/postprocess.py

@@ -32,6 +32,7 @@
     "watershed_split",
     "stitch_3d",
     "intersection_over_union",
+    "apply_binary_postprocessing",
 ]
 
 
@@ -248,3 +249,108 @@ def _label_overlap(x: np.ndarray, y: np.ndarray) -> np.ndarray:
     for i in range(len(x)):
         overlap[x[i], y[i]] += 1
     return overlap
+
+
+def apply_binary_postprocessing(
+    pred: np.ndarray, config: 'BinaryPostprocessingConfig'
+) -> np.ndarray:
+    """Apply binary segmentation postprocessing pipeline.
+
+    Pipeline order:
+        1. Threshold predictions to binary mask using threshold_range
+        2. Apply median filter (optional)
+        3. Apply morphological opening (erosion + dilation)
+        4. Apply morphological closing (dilation + erosion)
+        5. Extract connected components and filter by size/keep top-k
+
+    Args:
+        pred (numpy.ndarray): Predicted foreground probability in range [0, 1].
+                             Shape can be 2D (H, W) or 3D (D, H, W).
+        config (BinaryPostprocessingConfig): Configuration for postprocessing pipeline.
+
+    Returns:
+        numpy.ndarray: Postprocessed binary mask (same shape as input).
+                       Values: 0 (background) or 1 (foreground).
+
+    Example:
+        >>> from connectomics.config import BinaryPostprocessingConfig, ConnectedComponentsConfig
+        >>> config = BinaryPostprocessingConfig(
+        ...     enabled=True,
+        ...     threshold_range=(0.8, 1.0),
+        ...     opening_iterations=2,
+        ...     connected_components=ConnectedComponentsConfig(top_k=1)
+        ... )
+        >>> pred = np.random.rand(128, 128)  # Random probabilities
+        >>> binary_mask = apply_binary_postprocessing(pred, config)
+    """
+    if not config or not config.enabled:
+        # Just threshold at 0.5 if postprocessing is disabled
+        return (pred > 0.5).astype(np.uint8)
+
+    # Step 1: Threshold to binary using threshold_range
+    # Use the minimum threshold from the range
+    threshold = config.threshold_range[0]
+    binary = (pred > threshold).astype(np.uint8)
+
+    # Step 2: Apply median filter (optional noise reduction)
+    if config.median_filter_size is not None:
+        binary = ndimage.median_filter(binary, size=config.median_filter_size)
+
+    # Step 3: Morphological opening (erosion + dilation) - removes small objects
+    if config.opening_iterations > 0:
+        binary = ndimage.binary_opening(
+            binary, iterations=config.opening_iterations
+        ).astype(np.uint8)
+
+    # Step 4: Morphological closing (dilation + erosion) - fills small holes
+    if config.closing_iterations > 0:
+        binary = ndimage.binary_closing(
+            binary, iterations=config.closing_iterations
+        ).astype(np.uint8)
+
+    # Step 5: Connected components filtering
+    if config.connected_components is not None and config.connected_components.enabled:
+        cc_config = config.connected_components
+
+        # Extract connected components
+        connectivity = cc_config.connectivity
+        labels = cc3d.connected_components(binary, connectivity=connectivity)
+
+        # Get component sizes
+        component_sizes = np.bincount(labels.ravel())
+        # Skip background (label 0)
+        component_sizes[0] = 0
+
+        # Filter by minimum size
+        if cc_config.min_size > 0:
+            small_components = np.where(component_sizes < cc_config.min_size)[0]
+            for label_id in small_components:
+                labels[labels == label_id] = 0
+
+        # Keep only top-k largest components
+        if cc_config.top_k is not None and cc_config.top_k > 0:
+            # Get sizes (excluding background)
+            sizes = component_sizes[1:]  # Skip background
+            label_ids = np.arange(1, len(component_sizes))
+
+            if len(sizes) > cc_config.top_k:
+                # Get indices of top-k largest components
+                top_k_indices = np.argsort(sizes)[-cc_config.top_k:]
+                top_k_labels = label_ids[top_k_indices]
+
+                # Create mask keeping only top-k
+                keep_mask = np.zeros_like(labels, dtype=bool)
+                for label_id in top_k_labels:
+                    keep_mask |= (labels == label_id)
+
+                labels = keep_mask.astype(np.uint8)
+            else:
+                # Convert labels back to binary (0/1)
+                labels = (labels > 0).astype(np.uint8)
+        else:
+            # Convert labels back to binary (0/1)
+            labels = (labels > 0).astype(np.uint8)
+
+        binary = labels
+
+    return binary

connectomics/decoding/segmentation.py

@@ -43,6 +43,7 @@ def decorator(func):
 
 
 __all__ = [
+    "decode_binary_thresholding",
     "decode_binary_cc",
     "decode_binary_watershed",
     "decode_binary_contour_cc",
@@ -52,6 +53,80 @@ def decorator(func):
 ]
 
 
+def decode_binary_thresholding(
+    predictions: np.ndarray,
+    threshold_range: Tuple[float, float] = (0.8, 1.0),
+) -> np.ndarray:
+    r"""Convert binary foreground probability maps to binary mask via simple thresholding.
+
+    This is a lightweight decoding function that applies thresholding to convert
+    probability predictions to binary segmentation masks. Unlike instance segmentation
+    methods, this produces a semantic segmentation (no individual instance IDs).
+
+    The function uses the minimum threshold from threshold_range to binarize predictions.
+    This is useful for simple binary segmentation tasks where instance separation is not needed.
+
+    Args:
+        predictions (numpy.ndarray): foreground probability of shape :math:`(C, Z, Y, X)` or :math:`(C, Y, X)`.
+            The first channel (predictions[0]) is used as the foreground probability.
+            Values should be in range [0, 1] (normalized) or [0, 255] (uint8).
+        threshold_range (tuple): Tuple of (min_threshold, max_threshold) for binarization.
+            Only the minimum threshold is used. Values >= min_threshold become foreground (1).
+            Default: (0.8, 1.0)
+
+    Returns:
+        numpy.ndarray: Binary segmentation mask with shape matching input spatial dimensions.
+            Values: 0 (background) or 1 (foreground).
+            For 3D input: shape :math:`(Z, Y, X)`
+            For 2D input: shape :math:`(Y, X)`
+
+    Examples:
+        >>> # 3D predictions (normalized [0, 1])
+        >>> predictions = np.random.rand(2, 64, 128, 128)  # (C, Z, Y, X)
+        >>> binary_mask = decode_binary_thresholding(predictions, threshold_range=(0.8, 1.0))
+        >>> print(binary_mask.shape)  # (64, 128, 128)
+        >>> print(np.unique(binary_mask))  # [0, 1]
+
+        >>> # 3D predictions (uint8 [0, 255])
+        >>> predictions = np.random.randint(0, 256, (2, 64, 128, 128), dtype=np.uint8)
+        >>> binary_mask = decode_binary_thresholding(predictions, threshold_range=(0.8, 1.0))
+
+        >>> # 2D predictions
+        >>> predictions = np.random.rand(2, 512, 512)  # (C, Y, X)
+        >>> binary_mask = decode_binary_thresholding(predictions, threshold_range=(0.5, 1.0))
+        >>> print(binary_mask.shape)  # (512, 512)
+
+    Note:
+        - **Auto-detection of value range**: Automatically handles both normalized [0, 1]
+          and uint8 [0, 255] predictions
+        - **2D/3D support**: Works with both 2D (C, Y, X) and 3D (C, Z, Y, X) inputs
+        - **Channel 0 usage**: Uses first channel (predictions[0]) as foreground probability
+        - **Simple thresholding**: No morphological operations or connected components
+        - **Post-processing**: Use binary postprocessing config for refinement (opening/closing/CC filtering)
+
+    See Also:
+        - :func:`decode_binary_cc`: Binary threshold + connected components (instance segmentation)
+        - :func:`decode_binary_watershed`: Binary threshold + watershed (instance segmentation)
+        - :class:`connectomics.config.BinaryPostprocessingConfig`: For morphological refinement
+    """
+    # Extract foreground probability (first channel)
+    semantic = predictions[0]
+
+    # Auto-detect if predictions are in [0, 1] or [0, 255] range
+    max_value = np.max(semantic)
+    if max_value > 1.0:
+        # Assume uint8 range [0, 255]
+        threshold = threshold_range[0] * 255
+    else:
+        # Assume normalized range [0, 1]
+        threshold = threshold_range[0]
+
+    # Apply thresholding
+    binary_mask = (semantic > threshold).astype(np.uint8)
+
+    return binary_mask
+
+
 def decode_binary_cc(
     predictions: np.ndarray,
     foreground_threshold: float = 0.8,