From bdf5203904e2e65661474f16d57507e3250ebbe7 Mon Sep 17 00:00:00 2001
From: Pierre Tocquin <ptocquin@uliege.be>
Date: Tue, 3 Dec 2024 18:31:32 +0100
Subject: [PATCH] feat(process): enhance segmentation correction with
 additional parameters

Added new configuration options (`area_threshold`, `closing_kernel_size`, `aspect_ratio_min`, `aspect_ratio_max`) for improved segmentation correction in processing pipeline. These changes preserve true nodes of Ranvier by filtering based on cluster size and aspect ratio, and fill interruptions in nerve segments using morphological operations. Updated the function and documentation to reflect these enhancements. This improves accuracy and reliability of image segmentation results.
---
 process_pipeline.py | 153 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 148 insertions(+), 5 deletions(-)

diff --git a/process_pipeline.py b/process_pipeline.py
index af21d38..04065b8 100644
--- a/process_pipeline.py
+++ b/process_pipeline.py
@@ -9,6 +9,9 @@ Description: This script automates the processing of CZI images by converting th
              - Resize factor for scaling images (`resize_factor`)
              - Path to the Ilastik project file (`project`)
              - Path to the Ilastik executable (`ilastik`)
+             - Minimum size of red clusters to preserve as true nodes of Ranvier (`area_threshold`)
+             - Kernel size for filling interruptions in nerve segments (`closing_kernel_size`)
+             - Minimum and maximum aspect ratios to identify true nodes of Ranvier (`aspect_ratio_min`, `aspect_ratio_max`)
 
 Usage: python process_pipeline.py file1.czi file2.czi ... --config /path/to/config.yaml
 
@@ -18,13 +21,31 @@ Example configuration file (config.yaml):
     resize_factor: 0.25
     project: /path/to/ilastik_project.ilp
     ilastik: /path/to/ilastik_executable
+    area_threshold: 50
+    closing_kernel_size: 11
+    aspect_ratio_min: 1.5
+    aspect_ratio_max: 5
+
+Features:
+    - Assembles mosaics from CZI tiles into single TIFF files.
+    - Performs image segmentation using Ilastik in headless mode.
+    - Corrects segmentation issues in the segmented image:
+        * Preserves true nodes of Ranvier by filtering red clusters based on size and aspect ratio.
+        * Applies morphological operations to remove noise (e.g., thin structures).
+        * Fills interruptions in nerve segments caused by white gaps.
+    - Resizes both original and segmented images for downstream analysis.
+    - Removes intermediate files to optimize disk usage.
 
 Dependencies:
     - PyYAML: For reading the YAML configuration file.
+      Installation: pip install pyyaml
     - aicspylibczi: For processing CZI files.
+      Installation: pip install aicspylibczi
     - tifffile: For handling TIFF files.
-    - OpenCV: For image resizing.
-    - Ilastik: For segmentation in headless mode.
+      Installation: pip install tifffile
+    - OpenCV: For image resizing and morphological operations.
+      Installation: pip install opencv-python-headless
+    - Ilastik: For segmentation in headless mode (requires Ilastik installation).
 
 License: MIT License
 Copyright (c) 2024 Pierre Tocquin
@@ -53,7 +74,11 @@ def load_config(config_path):
     with open(config_path, 'r') as file:
         config = yaml.safe_load(file)
     
-    required_keys = ['output_dir', 'compression', 'resize_factor', 'project', 'ilastik']
+    required_keys = [
+        'output_dir', 'compression', 'resize_factor', 
+        'project', 'ilastik', 'area_threshold', 
+        'closing_kernel_size', 'aspect_ratio_min', 'aspect_ratio_max'
+    ]    
     for key in required_keys:
         if key not in config:
             raise ValueError(f"Missing required configuration key: {key}")
@@ -63,13 +88,32 @@ def load_config(config_path):
 
 def process_czi_file(file_path, config):
     """
-    Processes a CZI file, assembles mosaics, analyzes them with Ilastik, and resizes outputs.
+    Processes a CZI file, assembles mosaics, analyzes them with Ilastik, and corrects segmentation issues.
     Each series within the CZI file is processed sequentially.
 
+    Workflow:
+        1. Assemble mosaic from CZI tiles into a single TIFF file.
+        2. Process the mosaic with Ilastik for segmentation.
+        3. Correct segmentation issues in the segmented image:
+            - Preserve true nodes of Ranvier by filtering red clusters based on size and aspect ratio.
+            - Apply morphological operations to smooth and refine segmentation results.
+        4. Resize both the original mosaic and segmented outputs.
+        5. Remove intermediate files to optimize disk usage.
+
     Args:
         file_path (str): Path to the CZI file to process.
-        config (dict): Configuration options loaded from YAML.
+        config (dict): Configuration options loaded from YAML, including:
+            - output_dir (str): Directory for saving processed files.
+            - compression (str): Compression type for TIFF files (e.g., "lzw").
+            - resize_factor (float): Scaling factor for resizing images.
+            - project (str): Path to the Ilastik project file.
+            - ilastik (str): Path to the Ilastik executable.
+            - area_threshold (int): Minimum size for preserving red clusters as nodes of Ranvier.
+            - closing_kernel_size (int): Kernel size for filling interruptions in blue regions.
+            - aspect_ratio_min (float): Minimum aspect ratio for identifying nodes of Ranvier.
+            - aspect_ratio_max (float): Maximum aspect ratio for identifying nodes of Ranvier.
     """
+
     print(f"Processing file: {file_path}")
 
     # Load the CZI file
@@ -139,6 +183,23 @@ def process_czi_file(file_path, config):
         # Process the generated TIFF with Ilastik
         segmented_path = run_ilastik(output_filename, config['project'], config['ilastik'])
 
+        # Load segmented image for correction
+        segmented_image = tiff.imread(segmented_path)
+
+        # Correct segmentation problems (red clusters and interruptions)
+        corrected_image = correct_segmentation_issues(
+            segmented_image,
+            area_threshold=config['area_threshold'],
+            closing_kernel_size=config['closing_kernel_size'],
+            aspect_ratio_min=config['aspect_ratio_min'],
+            aspect_ratio_max=config['aspect_ratio_max']
+        )
+
+        # Save corrected image
+        tiff.imwrite(segmented_path, corrected_image.astype(np.uint8))
+        print(f"Corrected segmentation issues in segmented image: {segmented_path}")
+
+
         # Resize the original TIFF and the segmented output
         resize_image(output_filename, config['resize_factor'], cv2.INTER_AREA)
         resize_image(segmented_path, config['resize_factor'], cv2.INTER_NEAREST_EXACT)
@@ -194,6 +255,88 @@ def resize_image(image_path, scale_factor, interpolation):
     return resized_path
 
 
+def correct_segmentation_issues(image, area_threshold=50, closing_kernel_size=11, aspect_ratio_min=1.5, aspect_ratio_max=5):
+    """
+    Corrects segmentation issues in the processed image:
+    1. Preserves large red clusters (true nodes of Ranvier) based on area and shape.
+    2. Fills interruptions in nerve segments by closing small gaps in blue regions.
+
+    Args:
+        image (numpy.ndarray): The segmented image where white (1), blue (2), and red (3) represent classes.
+        area_threshold (int): Minimum area to preserve a red cluster as a true node of Ranvier.
+        closing_kernel_size (int): Size of the morphological closing kernel (must be odd).
+        aspect_ratio_min (float): Minimum aspect ratio to preserve a cluster as a true node of Ranvier.
+        aspect_ratio_max (float): Maximum aspect ratio to preserve a cluster as a true node of Ranvier.
+
+    Returns:
+        numpy.ndarray: Corrected segmented image.
+    """
+    # Step 1: Preserve large red clusters (true nodes of Ranvier)
+    red_mask = (image == 3).astype(np.uint8)
+    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(red_mask, connectivity=8)
+
+    true_red_clusters = set()
+    for label in range(1, num_labels):
+        area = stats[label, cv2.CC_STAT_AREA]
+        if area < area_threshold:
+            continue  # Ignore small clusters
+
+        # Extract the mask for the current cluster
+        cluster_mask = (labels == label).astype(np.uint8)
+
+        # Calculate bounding box and aspect ratio
+        contours, _ = cv2.findContours(cluster_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if not contours:
+            continue
+
+        # Use the largest contour for the cluster
+        contour = max(contours, key=cv2.contourArea)
+
+        # Calculate the oriented bounding box
+        rect = cv2.minAreaRect(contour)
+        width, height = rect[1]  # Dimensions of the bounding box
+
+        if width == 0 or height == 0:
+            continue  # Skip invalid clusters
+
+        # Calculate aspect ratio
+        aspect_ratio = max(width, height) / min(width, height)
+        if aspect_ratio_min <= aspect_ratio <= aspect_ratio_max:
+            true_red_clusters.add(label)
+
+    # Iterate over all red pixels to remove small red clusters
+    red_pixels = np.argwhere(image == 3)
+    for y, x in red_pixels:
+        if labels[y, x] not in true_red_clusters:
+            # Convert small red clusters to blue or white based on neighborhood
+            y1, y2 = max(0, y - 1), min(image.shape[0], y + 2)
+            x1, x2 = max(0, x - 1), min(image.shape[1], x + 2)
+            neighborhood = image[y1:y2, x1:x2]
+            blue_count = np.sum(neighborhood == 2)
+            white_count = np.sum(neighborhood == 1)
+
+            if blue_count > white_count:
+                image[y, x] = 2  # Convert to blue
+            else:
+                image[y, x] = 1  # Convert to white
+
+    # Step 2: Fill interruptions in blue regions (nerve segments)
+    blue_mask = (image == 2).astype(np.uint8)
+
+    # Apply morphological closing to fill small gaps
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (closing_kernel_size, closing_kernel_size))
+    closed_blue_mask = cv2.morphologyEx(blue_mask, cv2.MORPH_CLOSE, kernel)
+
+    # Identify newly filled areas (previously white and now part of the blue region)
+    newly_filled = (closed_blue_mask == 1) & (image == 1)
+    image[newly_filled] = 2  # Convert these pixels to blue
+
+    # Apply a median filter to smooth the final image
+    image = cv2.medianBlur(image, 7)
+
+    return image
+
+
 def main():
     parser = argparse.ArgumentParser(description="Automate CZI to Ilastik processing pipeline.")
     parser.add_argument("files", nargs='+', help="List of CZI files to process.")
-- 
GitLab