mask_segmentation_with_dapi.py

from skimage import filters, morphology, measure
from typing import Optional
import numpy as np

def apply_tissue_mask(dapi_image: np.ndarray, segmentation: np.ndarray, sigma: float = 2.0, 
                         area_threshold: int = 1000, dilation_radius: int = 3) -> np.ndarray:
    """
    Removes border nuclei from a DAPI-stained image by applying a mask based on segmentation.

    Parameters:
    - dapi_image (np.ndarray): Input grayscale DAPI image as a NumPy array.
    - segmentation (np.ndarray): Binary or labeled segmentation image to be masked.
    - sigma (float, optional): Sigma for Gaussian blur to smoothen the DAPI image. Default is 2.0.
    - area_threshold (int, optional): Minimum area for small hole removal. Default is 1000.
    - dilation_radius (int, optional): Radius for morphological dilation. Default is 3.
    
    Returns:
    - np.ndarray: Masked segmentation image with border nuclei removed.

    Raises:
    - ValueError: If input images are of different shapes or if they are not 2D or 3D arrays.
    """
    # Validate input dimensions
    if dapi_image.shape != segmentation.shape:
        raise ValueError("Input images must have the same shape.")

    # Step 1: Apply Gaussian blur to the DAPI image
    blurred_image = filters.gaussian(dapi_image, sigma=sigma)
    print(f"Applied Gaussian blur with sigma={sigma}. Shape: {blurred_image.shape}")

    # Step 2: Threshold the blurred image using Otsu's method
    threshold_value = filters.threshold_otsu(blurred_image)
    binary_image = blurred_image > threshold_value
    print(f"Applied Otsu thresholding. Threshold value: {threshold_value}")

    # Step 3: Fill holes in the binary image
    filled_image = morphology.remove_small_holes(binary_image, area_threshold=area_threshold)
    print(f"Filled small holes in binary image. Area threshold: {area_threshold}")

    # Step 4: Label connected regions in the filled image
    labeled_image = measure.label(filled_image)
    print(f"Labeled connected regions. Number of regions: {labeled_image.max()}")

    # Step 5: Merge small regions using morphological dilation
    merged_image = morphology.dilation(labeled_image, morphology.disk(dilation_radius))
    print(f"Applied morphological dilation with radius={dilation_radius}")

    # Step 6: Apply merged mask to the segmentation image
    masked_image = segmentation * (merged_image > 0)
    print("Applied merged mask to segmentation image.")

    return masked_image